aoc24/day16/d16p2.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
    let path: [Cell]
    var description: String { return "\(cell) from \(origin) @\(costSoFar)" }
}

struct CellPair : Hashable, CustomStringConvertible {
    let a: Cell
    let b: Cell
    var description: String { return "\(a),\(b)" }
}

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 ? "██" : "  " }
        }
        print("    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4")
        return 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]], [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) -> [Cell:(Int, [Cell])] {
    var q: [CellCost] = [(cell: start, cost: 0)]
    var visited: [Cell:(Int, [Cell])] = [:]
    while !q.isEmpty {
        let e = q.removeLast()
        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.append((cell: cell, cost: newCost))
            }
        }
    }
    return visited
}

func trace(_ visited: [Cell:(Int, [Cell])], 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] }
    }
}

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)")