d21p2: clean up code and properly calculate previously manually-determined best key sequences

day21/d21p2.swift

@@ -7,26 +7,15 @@ let (n7, n8, n9, n4, n5, n6, n1, n2, n3, n0) =
 let (nL, nR, nU, nD) = ("<", ">", "^", "v")
 
 protocol Pad {
-    static func seq(_ k1: Key, _ k2: Key) -> Set<[Key]>
+    static var pos: [Key: (Int, Int)] { get }
+    static var cache: [[Key]: Set<[Key]>] { get set }
+    static func valid(start: Key, seq: [Key]) -> Bool
 }
 
-struct KeyPair : Hashable {
+extension Pad {
-    let a: Key
+    // returns set of reasonable key sequences to move from k1 to k2
-    let b: Key
-}
-
-struct Numpad {
-    static let pos: [Key: (Int, Int)] = [
-        n7: (0, 0), n8: (0, 1), n9: (0, 2),
-        n4: (1, 0), n5: (1, 1), n6: (1, 2),
-        n1: (2, 0), n2: (2, 1), n3: (2, 2),
-                    n0: (3, 1), nA: (3, 2)
-    ]
-    static var cache: [KeyPair: Set<[Key]>] = [:]
-
     static func seq(_ k1: Key, _ k2: Key) -> Set<[Key]> {
-        let cacheKey = KeyPair(a: k1, b: k2)
+        if let cached = cache[[k1, k2]] {
-        if let cached = cache[cacheKey] {
             return cached
         }
         let (pos1, pos2) = (pos[k1]!, pos[k2]!)
@@ -37,17 +26,12 @@ struct Numpad {
                 Array(repeating: horKey, count: abs(pos1.1 - pos2.1)),
                 Array(repeating: horKey, count: abs(pos1.1 - pos2.1)) +
                 Array(repeating: verKey, count: abs(pos1.0 - pos2.0))
-            ].filter { seq in
+            ].filter { seq in valid(start: k1, seq: seq) }
-                (k1 != n1 || !seq.starts(with: [nD])) &&
+        cache[[k1, k2]] = Set(seqs)
-                (k1 != n4 || !seq.starts(with: [nD, nD])) &&
+        return cache[[k1, k2]]!
-                (k1 != n7 || !seq.starts(with: [nD, nD, nD])) &&
-                (k1 != n0 || !seq.starts(with: [nL])) &&
-                (k1 != nA || !seq.starts(with: [nL, nL]))
-            }
-        cache[cacheKey] = Set(seqs)
-        return cache[cacheKey]!
     }
 
+    // returns set of possible key sequences to press all |keys| from |start|
     static func press(_ start: Key, _ keys: ArraySlice<Key>) -> [[Key]] {
         if keys.count == 0 {
             return [[]]
@@ -60,64 +44,36 @@ struct Numpad {
     }
 }
 
-struct Dirpad {
+struct Numpad : Pad {
+    static let pos: [Key: (Int, Int)] = [
+        n7: (0, 0), n8: (0, 1), n9: (0, 2),
+        n4: (1, 0), n5: (1, 1), n6: (1, 2),
+        n1: (2, 0), n2: (2, 1), n3: (2, 2),
+                    n0: (3, 1), nA: (3, 2)
+    ]
+    static func valid(start: Key, seq: [Key]) -> Bool {
+        return
+            (start != n1 || !seq.starts(with: [nD])) &&
+            (start != n4 || !seq.starts(with: [nD, nD])) &&
+            (start != n7 || !seq.starts(with: [nD, nD, nD])) &&
+            (start != n0 || !seq.starts(with: [nL])) &&
+            (start != nA || !seq.starts(with: [nL, nL]))
+    }
+    static var cache: [[Key]: Set<[Key]>] = [:]
+}
+
+struct Dirpad : Pad {
     static let pos: [Key: (Int, Int)] = [
                     nU: (0, 1), nA: (0, 2),
         nL: (1, 0), nD: (1, 1), nR: (1, 2)
     ]
-    static var cache: [KeyPair: [Key]] = [
+    static func valid(start: Key, seq: [Key]) -> Bool {
-        // <<0 <v- <>- <^. <A.
+        return
-        // v<- vv0 v>- v^- vA
+            (start != nU || !seq.starts(with: [nL])) &&
-        // ><- >v- >>0 >^  >A-
+            (start != nA || !seq.starts(with: [nL, nL])) &&
-        // ^<. ^v- ^>  ^^0 ^A-
+            (start != nL || !seq.starts(with: [nU]))
-        // A<. Av  A>- A^- AA0
-        // Noops
-        KeyPair(a: nL, b: nL): [],
-        KeyPair(a: nR, b: nR): [],
-        KeyPair(a: nU, b: nU): [],
-        KeyPair(a: nD, b: nD): [],
-        KeyPair(a: nA, b: nA): [],
-        // Only solutions due to gap
-        KeyPair(a: nL, b: nU): [nR, nU],
-        KeyPair(a: nL, b: nA): [nR, nR, nU],
-        KeyPair(a: nU, b: nL): [nD, nL],
-        KeyPair(a: nA, b: nL): [nD, nL, nL],
-        // Only solutions
-        KeyPair(a: nL, b: nR): [nR, nR],
-        KeyPair(a: nR, b: nL): [nL, nL],
-        KeyPair(a: nL, b: nD): [nR],
-        KeyPair(a: nD, b: nL): [nL],
-        KeyPair(a: nU, b: nA): [nR],
-        KeyPair(a: nA, b: nU): [nL],
-        KeyPair(a: nD, b: nR): [nR],
-        KeyPair(a: nR, b: nD): [nL],
-        KeyPair(a: nA, b: nR): [nD],
-        KeyPair(a: nR, b: nA): [nU],
-        KeyPair(a: nU, b: nD): [nD],
-        KeyPair(a: nD, b: nU): [nU],
-        // Known short ones. TODO: proper DP this part instead of by hand
-        KeyPair(a: nA, b: nD): [nL, nD],
-        KeyPair(a: nD, b: nA): [nU, nR],
-        KeyPair(a: nR, b: nU): [nL, nU],
-        KeyPair(a: nU, b: nR): [nD, nR]
-    ]
-
-    static func seq(_ k1: Key, _ k2: Key) -> [Key] {
-        return cache[KeyPair(a: k1, b: k2)]!
     }
-
+    static var cache: [[Key]: Set<[Key]>] = [:]
-    static func press(_ start: Key, _ keys: ArraySlice<Key>) -> [Key] {
-        if keys.count == 0 {
-            return []
-        }
-        return seq(start, keys[keys.startIndex]) + [nA] +
-            press(keys[keys.startIndex], keys.dropFirst())
-    }
-}
-
-func readInput(_ filePath: String) throws -> [[Key]] {
-    return try String(contentsOfFile: filePath, encoding: .ascii)
-        .split(separator: "\n").map { Array($0).map(String.init) }
 }
 
 extension Array {
@@ -133,46 +89,54 @@ extension Array {
     }
 }
 
-struct LKs : Hashable {
+struct LKs : Hashable, CustomStringConvertible {
     let l: Int
     let ks: [Key]
+    var description: String { return "(\(l),\(ks.joined()))" }
 }
 
-// Returns a map of a sequence of pressed starting at A and ending at A
+// Returns a lookup table from a sequence of pressed ending at A
-// to how many key presses total for that
+// to how many key presses total for that sequence
-func bestLevel1() -> [LKs: Int] {
+func best(maxLevel: Int) -> [LKs: Int] {
     var res: [LKs: Int] = [:]
     [nL, nR, nU, nD].combos(maxl: 3).forEach { seq in
-        res[LKs(l: 1, ks: seq)] = Dirpad.press(nA, seq + [nA]).count
+        let len = Dirpad.press(nA, seq + [nA]).map { $0.count}.min()!
+        res[LKs(l: 1, ks: seq)] = len
     }
-    return res
+    var level = 2
-}
+    while level <= maxLevel {
-func best(maxLevel: Int) -> [LKs: Int] {
+        let maxl = (level == maxLevel) ? 5 : 3
-    var res = bestLevel1()
+        [nL, nR, nU, nD].combos(maxl: maxl).forEach { seq in
-    if maxLevel == 1 { return res }
-    (2...maxLevel).forEach { level in
-        [nL, nR, nU, nD].combos(maxl: 3).forEach { seq in
             res[LKs(l: level, ks: Array(seq))] = Dirpad.press(nA, seq + [nA])
-                .split(separator: "A", omittingEmptySubsequences: false).dropLast()
+                .map { ks in
-                .map { res[LKs(l: level-1, ks: Array($0))]! }.reduce(0, +)
+                    ks.split(separator: "A", omittingEmptySubsequences: false)
+                        .dropLast()
+                        .map { res[LKs(l: level-1, ks: Array($0))]! }
+                        .reduce(0, +)
+                }.min()!
         }
+        level += 1
     }
     return res
 }
 
-let levels = Int(CommandLine.arguments[2])! - 1
+func readInput(_ filePath: String) throws -> [[Key]] {
+    return try String(contentsOfFile: filePath, encoding: .ascii)
+        .split(separator: "\n").map { Array($0).map(String.init) }
+}
+
+let levels = Int(CommandLine.arguments[2])!
 let bests = best(maxLevel: levels)
 
 var answer = 0
 for passwd in try readInput(CommandLine.arguments[1]) {
-    let n = Numpad.press(nA, passwd[...])
+    let cost = Numpad.press(nA, passwd[...]).map { ks in
-        .map { Dirpad.press(nA, $0[...]) }.map { ks in
+            ks.split(separator: "A", omittingEmptySubsequences: false)
-            return ks
+                .dropLast()
-                .split(separator: "A", omittingEmptySubsequences: false).dropLast()
                 .map { bests[LKs(l: levels, ks: Array($0))]! }
                 .reduce(0, +)
         }.min()!
-    answer += n * Int(passwd.joined().dropLast())!
+    answer += cost * Int(passwd.joined().dropLast())!
 }
 print(answer)