p18.swift

import Foundation

struct Position: Hashable {
    let x: Int
    let y: Int
    var adjacent: [Position] {
        return [Position(x: x-1, y: y),
                Position(x: x+1, y: y),
                Position(x: x, y: y-1),
                Position(x: x, y: y+1)]
    }
}

enum Space: CustomStringConvertible {
    case wall
    case open
    case key(Character)
    case door(Character)
    case start
    init(_ s: Character) {
        switch s {
        case "#":   self = .wall
        case ".":   self = .open
        case "@":   self = .start
        case let a: self = a.isLowercase ? .key(a) : .door(a)
        }
    }
    var description: String {
        switch self {
        case .wall: return "#"
        case .start: return "@"
        case .open: return "."
        case .key(let s): return String(s)
        case .door(let s): return String(s)
        }
    }
}

typealias Graph = [Character:[Character:Int]]

struct State {
    let graph: Graph
    let collected: [Character]
    let distanceTravelled: Int
    let paths: [[Character:Int]] // One for each robot
    
    var collectedInOrder: String {
        guard let l = collected.last else { return "" }
        return String(collected.dropLast().sorted() + [l])
    }
    
    var collectedKeys: [Character] { return collected.filter({ $0.isLowercase }) }
    
    init(string s: String) {
        let rows = s.components(separatedBy: .newlines)
        let height = rows.count
        let width = rows.reduce(0, { max($0, $1.count) })
        var g = rows.map({ (s: String) -> [Space] in
            var a = Array<Space>(repeating: .open, count: width)
            for i in s.enumerated() {
                a[i.offset] = Space(i.element)
            }
            return a
        })
        var start = [Position]()
        var nodes = [Character:Position]()
        for r in 0..<height { for c in 0..<width {
            switch g[r][c] {
            case .open, .wall:                  break
            case .start:                        start.append(Position(x: c, y: r))
            case .door(let a), .key(let a):     nodes[a] = Position(x: c, y: r)
            }
        } }

        func distanceBetween(p1: Position, p2: Position) -> Int? {
            // BFS search in grid "g" from p1 to p2.
            // Stop at both keys and doors.
            var steps = 0
            var leafStates = [p1]
            var oldStates = [Position]()
            while true {
                steps = steps + 1
                var additionalStates = [Position]()
                for leaf in leafStates {
                    for move in leaf.adjacent {
                        if move == p2 { return steps }
                        switch g[move.y][move.x] {
                        case .wall, .door, .key: continue
                        case .open, .start: break
                        }
                        if oldStates.contains(move) { continue }
                        additionalStates.append(move)
                    }
                }
                oldStates = leafStates
                leafStates = additionalStates
                if additionalStates.count == 0 {
                    return nil
                }
            }
        }
        
        var out = Graph()
        for n in nodes {
            var distances = [Character:Int]()
            for m in nodes {
                if m == n { continue }
                if let d = distanceBetween(p1: n.value, p2: m.value) {
                    distances[m.key] = d
                }
            }
            out[n.key] = distances
        }
        
        self.paths = start.map( { (position: Position) -> [Character:Int] in
            let distances: [(Character, Int)] = nodes.compactMap({
                if let d = distanceBetween(p1: position, p2: $0.value) {
                    return ($0.key, d)
                } else {
                    return nil
                }
            })
            return  [Character:Int](uniqueKeysWithValues: distances)
        })
        self.distanceTravelled = 0
        self.graph = out
        self.collected = []
    }
    
    init(oldState: State, visiting: Character, distance fromLastToVisited: Int, robot: Int) {
        self.distanceTravelled = fromLastToVisited + oldState.distanceTravelled
        self.collected = oldState.collected + [visiting]

        // Get a list of nodes connected to visiting, and their distance.
        var distanceToVisiting = [Character:Int]()
        for node in oldState.graph {
            if let distance = node.value[visiting] {
                distanceToVisiting[node.key] = distance
            }
        }
            
        var graph = Graph()
        for nodeIterator in oldState.graph {
            let (node, connections) = nodeIterator
            if node == visiting { continue }
            if connections[visiting] == nil {
                // Untouched by removing visiting
                graph[node] = connections
                continue
            }
            
            // This node is connected to visiting.
            var d = connections
            // The distance to the visited node should be removed.
            guard let m = d.removeValue(forKey: visiting) else {
                fatalError("Not connected to visiting?")
            }

            for v in distanceToVisiting {
                let (otherNode, distanceFromOtherNodeToVisiting) = v
                if otherNode == node { continue }
                if let oldDistance = d[otherNode], oldDistance > distanceFromOtherNodeToVisiting + m {
                    d[otherNode] = distanceFromOtherNodeToVisiting + m
                }
                if d[otherNode] == nil {
                    d[otherNode] = distanceFromOtherNodeToVisiting + m
                }
            }
            graph[node] = d
        }
        var paths = oldState.paths
        paths[robot] = oldState.graph[visiting]!
        self.paths = paths
        self.graph = graph
    }
    
    var possibleNewStates: [State] {
        var s = [State]()
        for p in self.paths.enumerated() {
            let availableStatesFromThisRobot: [State] = p.element.filter({ $0.key.isLowercase || collected.contains(Character($0.key.lowercased())) }).map({ State(oldState: self, visiting: $0.key, distance: $0.value, robot: p.offset) })
            s.append(contentsOf: availableStatesFromThisRobot)
        }
        return s
    }
}

let path = CommandLine.arguments[1]
let input = try! String(contentsOf: URL(fileURLWithPath: path))
let initial = State(string: input)
let numberOfKeysToCollect = initial.graph.keys.filter({ $0.isLowercase }).count

// Start with a DFS and always choose the shortest leaf, to get a
// maximum upper bound on the length
func depthFirstSearch(from state: State) -> Int? {
    guard let available = state.possibleNewStates.sorted(by: { $0.distanceTravelled < $1.distanceTravelled }).first else { return nil }
    
    if available.collectedKeys.count == numberOfKeysToCollect {
        return available.distanceTravelled
    }
    return depthFirstSearch(from: available)
}

var totalLength = depthFirstSearch(from: initial) ?? Int.max
print("Upper bound on number of steps: \(totalLength)")
    
// Then do a BFS,
var leafStates = initial.possibleNewStates
while leafStates.count > 0 {
    var newLeafStates = [State]()
    for leaf in leafStates {
        let p = leaf.possibleNewStates.filter({ $0.distanceTravelled < totalLength })
        newLeafStates.append(contentsOf: p)
    }
    // Important: purge equal states - being at the same node n and having the same 0..(n-1) nodes
    // but in different order. Keep the shorest equivalent state.
    var shortestPaths = [String:State]()
    for leaf in newLeafStates {
        let s = leaf.collectedInOrder
        if let j = shortestPaths[s], j.distanceTravelled < leaf.distanceTravelled {} else {
            shortestPaths[s] = leaf
        }
    }
    leafStates = [State](shortestPaths.values)
    
    let lengthOfCompletedNewStates = leafStates.filter({ $0.collectedKeys.count == numberOfKeysToCollect }).map({ $0.distanceTravelled})
    totalLength = min(totalLength, lengthOfCompletedNewStates.min() ?? Int.max)
}
print("Number of steps to collect all keys: \(totalLength)")