aoc24/day21/d21p2.swift

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 var pos: [Key: (Int, Int)] { get }
    static var cache: [[Key]: Set<[Key]>] { get set }
    static func valid(start: Key, seq: [Key]) -> Bool
}

extension Pad {
    // returns set of reasonable key sequences to move from k1 to k2
    static func seq(_ k1: Key, _ k2: Key) -> Set<[Key]> {
        if let cached = cache[[k1, k2]] {
            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 valid(start: k1, seq: seq) }
        cache[[k1, k2]] = Set(seqs)
        return cache[[k1, k2]]!
    }

    // 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 [[]]
        }
        return seq(start, keys[keys.startIndex]).flatMap { firstSeq in
                press(keys[keys.startIndex], keys.dropFirst()).map { restSeq in
                    firstSeq + [nA] + restSeq
                }
            }
    }
}

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 func valid(start: Key, seq: [Key]) -> Bool {
        return
            (start != nU || !seq.starts(with: [nL])) &&
            (start != nA || !seq.starts(with: [nL, nL])) &&
            (start != nL || !seq.starts(with: [nU]))
    }
    static var cache: [[Key]: Set<[Key]>] = [:]
}

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, CustomStringConvertible {
    let l: Int
    let ks: [Key]
    var description: String { return "(\(l),\(ks.joined()))" }
}

// Returns a lookup table from a sequence of pressed ending at A
// to how many key presses total for that sequence
func best(maxLevel: Int) -> [LKs: Int] {
    var res: [LKs: Int] = [:]
    [nL, nR, nU, nD].combos(maxl: 3).forEach { seq in
        let len = Dirpad.press(nA, seq + [nA]).map { $0.count}.min()!
        res[LKs(l: 1, ks: seq)] = len
    }
    var level = 2
    while level <= maxLevel {
        let maxl = (level == maxLevel) ? 5 : 3
        [nL, nR, nU, nD].combos(maxl: maxl).forEach { seq in
            res[LKs(l: level, ks: Array(seq))] = Dirpad.press(nA, seq + [nA])
                .map { ks in
                    ks.split(separator: "A", omittingEmptySubsequences: false)
                        .dropLast()
                        .map { res[LKs(l: level-1, ks: Array($0))]! }
                        .reduce(0, +)
                }.min()!
        }
        level += 1
    }
    return res
}

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 cost = Numpad.press(nA, passwd[...]).map { ks in
            ks.split(separator: "A", omittingEmptySubsequences: false)
                .dropLast()
                .map { bests[LKs(l: levels, ks: Array($0))]! }
                .reduce(0, +)
        }.min()!
    answer += cost * Int(passwd.joined().dropLast())!
}
print(answer)