aoc24/day16/d16p1.swift

import Foundation

typealias CellCost = (cell: Cell, cost: Int)

enum Dir : String, CaseIterable, CustomStringConvertible {
    case n = "▲"
    case s = "▼"
    case e = "▶"
    case w = "◀"
    var description: String { return self.rawValue }
    func turnCost(to dir: Self) -> Int? {
        if self == dir { return 0 }
        if (self == .n && dir == .s) || (self == .s && dir == .n) ||
           (self == .e && dir == .w) || (self == .w && dir == .e) { return nil }
        return 1000
    }
}

struct Cell : Hashable, CustomStringConvertible {
    let (i, j): (Int, Int)
    let dir: Dir
    var description: String { return "(\(i),\(j))\(dir)" }
    func neighbor(_ dir: Dir) -> CellCost? {
        if let cost = self.dir.turnCost(to: dir) {
            let c = cost + 1
            switch dir {
            case Dir.n: return (cell: Cell(i: i-1, j: j, dir: dir), cost: c)
            case Dir.s: return (cell: Cell(i: i+1, j: j, dir: dir), cost: c)
            case Dir.e: return (cell: Cell(i: i, j: j+1, dir: dir), cost: c)
            case Dir.w: return (cell: Cell(i: i, j: j-1, dir: dir), cost: c)
            }
        } else {
            return nil
        }
    }
}

struct QueueEntry : Hashable, CustomStringConvertible {
    let cell: Cell
    let origin: Cell
    let costSoFar: Int
    var description: String { return "\(cell) from \(origin) @\(costSoFar)" }
}

struct CellPair : Hashable {
    let a: Cell
    let b: Cell
}

struct Maze : CustomStringConvertible {
    let walls: [[Bool]]
    let (w, h): (Int, Int)
    var start: Cell { return Cell(i: h-2, j: 1, dir: Dir.e) }
    var ends: Set<Cell> {
        return Set(Dir.allCases.map { Cell(i: 1, j: w-2, dir: $0) })
    }
    var description: String {
        let s: [[String]] = walls.enumerated().map { i, r in
            [String(format: "%3d ", i)] + r.map { $0 ? "██" : "  " }
        }
        return ("    " + (0..<w).map { " \($0 % 10)" }.joined()) +
            s.map { $0.joined() }.joined(separator: "\n")
    }
    init(fromFile f: String) throws {
        let content = try String(contentsOfFile: f, encoding: .ascii)
        let lines = content.split(separator: "\n")
        (h, w) = (lines.count, lines[0].count)
        walls = lines.map { line in line.map { cell in cell == "#" } }
    }

    func neighbors(of cell: Cell) -> [CellCost] {
        return Dir.allCases
            .compactMap { cell.neighbor($0) }
            .filter { nb in !walls[nb.cell.i][nb.cell.j] }
    }

    func graph() -> [Cell:[CellCost]] {
        var edges: [Cell:[(Cell, Int)]] = [start: []]
        var q: [QueueEntry] = [
            QueueEntry(cell: start, origin: start, costSoFar: 0)
        ]
        var seen: Set<CellPair> = []
        while !q.isEmpty {
            let e = q.removeLast()
            var origin = e.origin
            var costSoFar = e.costSoFar
            let nbs = neighbors(of: e.cell)
            if ends.contains(e.cell) || (nbs.count > 1 && e.cell != start) {
                if seen.contains(CellPair(a: origin, b: e.cell)) {
                    continue
                }
                let upd = edges[origin, default: []] + [(e.cell, costSoFar)]
                edges[origin] = upd
                seen.insert(CellPair(a: origin, b: e.cell))
                origin = e.cell
                costSoFar = 0
            }
            if !ends.contains(e.cell) {
                nbs.map { cell, cost in
                    QueueEntry(
                        cell: cell, origin: origin, costSoFar: costSoFar + cost
                    )
                }.forEach { q.append($0) }
            }
        }
        ends.forEach { edges[$0] = [] }
        return edges
    }

    func search() -> Int {
        let edges = graph()
        // edges.forEach { k, v in print(" - \(k): \(v)") }

        var q = Heap<CellCost>(comparator: { $0.cost < $1.cost })
        q.insert((cell: start, cost: 0))
        var visited: [Cell:Int] = [:]
        while let e = q.pop() {
            if let node = edges[e.cell] {
                for (cell, cost) in node {
                    if let prevCost = visited[cell] {
                        if e.cost + cost >= prevCost {
                            continue
                        }
                    }
                    visited[cell] = e.cost + cost
                    q.insert((cell: cell, cost: e.cost + cost))
                }
            }
        }
        return ends.compactMap { visited[$0] }.min()!
    }
}

@main
struct AoC {
    static func main() throws {
        let maze = try Maze(fromFile: CommandLine.arguments[1])
        print(maze)
        print("minCost: \(maze.search())")
    }
}