Solve day07 part 2 and generalize CrabAlignment struct

Author: maximiliantech

2021.playground/Pages/Day07.xcplaygroundpage/Contents.swift

@@ -7,32 +7,55 @@ struct CrabAlignment {
     var position: Int
     var horizontalPositions: [Int]
 
-    init(for position: Int, horizontalPositions: [Int]) {
+    init(for position: Int, horizontalPositions: [Int], constantRate: Bool) {
         self.position = position
         self.horizontalPositions = horizontalPositions
-        var total = 0
-        for curr_pos in horizontalPositions {
-            total += abs(curr_pos - position)
+        var total: Int = 0
+        if constantRate {
+            for curr_pos in horizontalPositions {
+                total += abs(curr_pos - position)
+            }
+        } else {
+            for curr_pos in horizontalPositions {
+                total += ( ( abs(curr_pos - position) * (abs(curr_pos - position) + 1) ) / 2 )
+            }
         }
         self.totalFuel = total
     }
 }
 
 func partOne() -> String {
     // load input and remove \n
-    let input = load(file: "input07_test", ofType: .txt)?.replacingOccurrences(of: "\n", with: "")
+    let input = load(file: "input07", ofType: .txt)?.replacingOccurrences(of: "\n", with: "")
     // turn numbers in [Int]
     let horizontalPositions = input?.split(separator: ",").map { Int($0)! }
 
     var crabAlignmentArray: [CrabAlignment] = []
-    for align_pos in horizontalPositions! {
-        crabAlignmentArray.append(CrabAlignment(for: align_pos, horizontalPositions: horizontalPositions!))
+    // Crabs can be aligned at any horizontal position
+    for align_pos in horizontalPositions!.min()! ... horizontalPositions!.max()! {
+        crabAlignmentArray.append(CrabAlignment(for: align_pos, horizontalPositions: horizontalPositions!, constantRate: true))
     }
-    // sort 
+    // sort array for the least total fuel
     crabAlignmentArray.sort { $0.totalFuel < $1.totalFuel }
     return crabAlignmentArray[0].totalFuel.description
 }
 
-print(partOne())
+func partTwo() -> String {
+    // load input and remove \n
+    let input = load(file: "input07", ofType: .txt)?.replacingOccurrences(of: "\n", with: "")
+    // turn numbers in [Int]
+    let horizontalPositions = input?.split(separator: ",").map { Int($0)! }
+    
+    var crabAlignmentArray: [CrabAlignment] = []
+    // Crabs can be aligned at any horizontal position
+    for align_pos in horizontalPositions!.min()! ... horizontalPositions!.max()! {
+        crabAlignmentArray.append(CrabAlignment(for: align_pos, horizontalPositions: horizontalPositions!, constantRate: false))
+    }
+    // sort array for the least total fuel
+    crabAlignmentArray.sort { $0.totalFuel < $1.totalFuel }
+    return crabAlignmentArray[0].totalFuel.description
+}
 
+print(partOne())
+print(partTwo())
 //: [Next](@next)