import Foundation typealias Key = String let (nA, nEmpty) = ("A", "") let (n7, n8, n9, n4, n5, n6, n1, n2, n3, n0) = ("7", "8", "9", "4", "5", "6", "1", "2", "3", "0") let (nL, nR, nU, nD) = ("<", ">", "^", "v") protocol Pad { static func seq(_ k1: Key, _ k2: Key) -> Set<[Key]> } struct KeyPair : Hashable { let a: Key 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[cacheKey] { return cached } let (pos1, pos2) = (pos[k1]!, pos[k2]!) let verKey = pos1.0 < pos2.0 ? nD : nU let horKey = pos1.1 < pos2.1 ? nR : nL let seqs: [[Key]] = [ // ??? Array(repeating: verKey, count: abs(pos1.0 - pos2.0)) + 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 (k1 != n1 || !seq.starts(with: [nD])) && (k1 != n4 || !seq.starts(with: [nD, nD])) && (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]! } static func press(_ start: Key, _ keys: ArraySlice) -> [[Key]] { if keys.count == 0 { return [[]] } return seq(start, keys[keys.startIndex]).flatMap { firstSeq in press(keys[keys.startIndex], keys.dropFirst()).map { restSeq in firstSeq + [nA] + restSeq } } } } struct Dirpad { 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]] = [ // <<0 - <^. - v^- vA // ><- >v- >>0 >^ >A- // ^<. ^v- ^> ^^0 ^A- // 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 func press(_ start: Key, _ keys: ArraySlice) -> [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 { func combos(l: Int) -> [[Element]] { guard count > 0 else { return [[]] } if l == 0 { return [[]] } return flatMap { e in combos(l: l - 1).map { $0 + [e] } } } func combos(maxl: Int) -> [[Element]] { return (0...maxl).flatMap { combos(l: $0) } } } struct LKs : Hashable { let l: Int let ks: [Key] } // Returns a map of a sequence of pressed starting at A and ending at A // to how many key presses total for that func bestLevel1() -> [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 } return res } func best(maxLevel: Int) -> [LKs: Int] { var res = bestLevel1() 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 { res[LKs(l: level-1, ks: Array($0))]! }.reduce(0, +) } } return res } let levels = Int(CommandLine.arguments[2])! - 1 let bests = best(maxLevel: levels) var answer = 0 for passwd in try readInput(CommandLine.arguments[1]) { let n = Numpad.press(nA, passwd[...]) .map { Dirpad.press(nA, $0[...]) }.map { ks in return ks .split(separator: "A", omittingEmptySubsequences: false).dropLast() .map { bests[LKs(l: levels, ks: Array($0))]! } .reduce(0, +) }.min()! answer += n * Int(passwd.joined().dropLast())! } print(answer)