Caching neighbours

Sources/09/09.swift

@@ -9,7 +9,7 @@ import Foundation
 import Runner
 import Collections
 
-struct Coord {
+struct Coord: Hashable {
     let row: Int
     let col: Int
     
@@ -33,31 +33,34 @@ struct Coord {
 }
 
 struct Field:Sequence, IteratorProtocol {
+    typealias Element = Coord
     let numbers: [[Int]]
     let numCols: Int
     let numRows: Int
     
+    var neighbors = [Coord:[Coord]]()
+    
     var current: Coord? = Coord(row: 0, col: 0)
     
     init(numbers: [[Int]]) {
         self.numbers = numbers
-        numCols = numbers.first!.count
+        numCols = numbers.first?.count ?? 0
         numRows = numbers.count
     }
     
-    func getNeighborsCoordsFor(coord: Coord) -> [Coord] {
+    mutating func getNeighborsCoordsFor(coord: Coord) -> [Coord] {
-        return [coord.getAbove(), coord.getRight(), coord.getBelow(), coord.getLeft()]
+        let n = neighbors[coord] ?? [coord.getAbove(), coord.getRight(), coord.getBelow(), coord.getLeft()]
             .compactMap { $0 }
             .filter { $0.row < numRows && $0.col < numCols }
+        neighbors[coord] = n
+        return n
     }
     
-    func lowPointValueAt(coord: Coord) -> Int? {
+    mutating func isLowPointAt(coord: Coord) -> Bool {
-        let value = numbers[coord.row][coord.col]
+        let value = self[coord]
         return getNeighborsCoordsFor(coord: coord)
             .map { numbers[$0.row][$0.col] }
             .allSatisfy { $0 > value }
-        ? value
-        : nil
     }
     
     private func getNextCoord() -> Coord? {
@@ -79,19 +82,59 @@ struct Field:Sequence, IteratorProtocol {
         }
         return current
     }
+    
+    subscript(index: Coord) -> Int {
+        get {
+            numbers[index.row][index.col]
+        }/*
+        set(newValue) {
+            // Perform a suitable setting action here.
+        }*/
+    }
 }
 
-struct Day09: Runnable {
+class  Day09: Runnable {
     let inputPath: String
+    var lowestNeighbors = [Coord: Coord]()
+    var lowPoints = Set<Coord>()
+    var field = Field(numbers: [[Int]]())
     
-    func run() {
+    required init(inputPath: String) {
+        self.inputPath = inputPath
+    }
+    
+    public func run() {
         let input = try! String(contentsOfFile: inputPath)
         let entries = input
             .trimmingCharacters(in: .newlines)
             .components(separatedBy: .newlines)
             .map { line in line.map { $0.wholeNumberValue! } }
-        let field = Field(numbers: entries)
+        field = Field(numbers: entries)
-        let sum = field.compactMap { field.lowPointValueAt(coord: $0) }.reduce(0, { $0 + $1 + 1 })
+        lowPoints = Set(field.filter { field.isLowPointAt(coord: $0) })
+        
+        //Part 1
+        let sum = lowPoints.map { field[$0] }.reduce(0, +) + lowPoints.count
         print(sum)
+        
+        // Part 2:
+        var basinSize = [Coord: Int]()
+        for coord in field {
+            guard field[coord] != 9 else {
+                continue
+            }
+            let bl = basinLocation(for: coord)
+            basinSize[bl, default: 0] += 1
+        }
+        let result = basinSize.values.sorted(by: >)[0..<3].reduce(1, *)
+        print(result)
+    }
+    
+    private func basinLocation(for coord: Coord) -> Coord {
+        guard !lowPoints.contains(coord) else {
+            return coord
+        }
+        let closestNeighbor = lowestNeighbors[coord] ?? field.getNeighborsCoordsFor(coord: coord).min(by: { field[$0] < field[$1] })!
+        lowestNeighbors[coord] = closestNeighbor
+        return basinLocation(for: closestNeighbor)
     }
 }

Sources/09/main.swift

@@ -8,5 +8,5 @@
 import Foundation
 import Runner
 
-Runner(target: Day09.self, day: "09", isTest: true).run()
+//Runner(target: Day09.self, day: "09", isTest: true).run()
-//Runner(target: Day09.self, day: "09", isTest: false).run()
+Runner(target: Day09.self, day: "09", isTest: false).run()