aoc24/day16/d16p2.swift

import Foundation

typealias CellCost = (cell: Cell, cost: Int)
typealias PathCostFromStart = [Cell:(Int, [Cell])]

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 {
    let cell: Cell
    let origin: Cell
    let costSoFar: Int
    let path: [Cell]
}

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

struct Maze {
    let walls: [[Bool]]
    let (w, h): (Int, Int)
    let start: Cell
    let ends: Set<Cell>
    init(fromFile f: String) throws {
        let content = try String(contentsOfFile: f, encoding: .ascii)
        let lines = content.split(separator: "\n")
        let (h, w) = (lines.count, lines[0].count)
        (self.h, self.w) = (h, w)
        walls = lines.map { line in line.map { cell in cell == "#" } }
        start = Cell(i: h-2, j: 1, dir: Dir.e)
        ends = Set(Dir.allCases.map { Cell(i: 1, j: w-2, dir: $0) })
    }

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

}

func search(graph edges: [Cell:[CellCost]], start: Cell) -> PathCostFromStart {
    var q = Heap<CellCost>(comparator: { l, r in l.cost < r.cost })
    q.insert((cell: start, cost: 0))
    var visited: PathCostFromStart = [:]
    while let e = q.pop() {
        if let node = edges[e.cell] {
            for (cell, cost) in node {
                let newCost = e.cost + cost
                if let (prevCost, prevCells) = visited[cell] {
                    if newCost > prevCost {
                        continue
                    }
                    if newCost == prevCost {
                        visited[cell] = (newCost, prevCells + [e.cell])
                        continue
                    }
                }
                visited[cell] = (newCost, [e.cell])
                q.insert((cell: cell, cost: newCost))
            }
        }
    }
    return visited
}

func trace(_ visited: PathCostFromStart, from: Cell, to: Cell) -> [[Cell]] {
    if from == to {
        return [[from]]
    }
    return visited[from, default: (-1, [])].1.flatMap {
        trace(visited, from: $0, to: to).map { $0 + [from] }
    }
}

@main
struct AoC {
    static func main() throws {
        let maze = try Maze(fromFile: CommandLine.arguments[1])
        let (graph, pathPairs) = maze.graph()
        let visited = search(graph: graph, start: maze.start)
        let minCost = maze.ends.compactMap { visited[$0] }.map { $0.0 }.min()!
        print("minCost: \(minCost)")
        let seats = maze.ends
            .filter { visited[$0, default: (0, [])].0 == minCost }
            .flatMap { trace(visited, from: $0, to: maze.start) }
            .map { zip($0, $0.dropFirst()) }
            .map { $0.flatMap { pathPairs[CellPair(a: $0.0, b: $0.1)]! } }
            .map { $0.map { cell in Cell(i: cell.i, j: cell.j, dir: Dir.n) } }
            .flatMap { $0 }
        print("seats: \(Set(seats).count)")
    }
}