aoc2021/Sources/15/15.swift

//
//  File.swift
//
//
//  Created by Max Nuding on 15s.12.21.
//

import Foundation
import Runner
import Collections

typealias ThreatLevel = Int

struct Coord: Hashable, CustomStringConvertible {
    let row: Int
    let col: Int
    
    var description: String { "(\(row), \(col))" }
    
    func neighbors() -> [Coord] {
        [
            Coord(row: row-1, col: col),
            Coord(row: row, col: col+1),
            Coord(row: row+1, col: col),
            Coord(row: row, col: col-1)
        ]
    }
}

class Day15: Runnable {
    let inputPath: String
    var lastRow: Int = 0
    var lastCol: Int = 0
    var parts = [[ThreatLevel]]()
    var neighbors = [Coord: [Coord]]()
    
    required init(inputPath: String) {
        self.inputPath = inputPath
    }
    
    private func expandMap() {
        var tmpParts = parts
        for (rowNum, row) in tmpParts.enumerated() {
            let newRow = [Int](0...4).flatMap { cr in
                row.map { rv in cr + rv }.map { $0 > 9 ? $0 - 9 : $0 }
            }
            parts[rowNum] = newRow
        }
        
        tmpParts = parts
        for num in 1...4 {
            for i in tmpParts.indices {
                let tr = tmpParts[i]
                let ntr = tr.map { $0 + num }.map { $0 > 9 ? $0 - 9 : $0 }
                parts.append(ntr)
            }
        }
        
        lastRow = parts.count - 1
        lastCol = (parts.first?.count ?? 0) - 1
    }
    
    public func run() {
        let input = try! String(contentsOfFile: inputPath)
        parts = input
            .trimmingCharacters(in: .newlines)
            .components(separatedBy: .newlines)
            .map { Array($0) }
            .map { ar in ar.map {c in ThreatLevel(c.description)!} }
        lastRow = parts.count - 1
        lastCol = (parts.first?.count ?? 0) - 1
        
        calculateSafestPath()
        expandMap()
        neighbors = [Coord: [Coord]]()
        calculateSafestPath()
    }
    
    private func calculateSafestPath() {
        // Distances to startNode
        var totalThreatLevel = [Coord: ThreatLevel]()
        
        var unvisited = Set(parts.enumerated().flatMap {row in
            row.element.indices.map { col in
                Coord(row: row.offset, col: col)
            }
        })
        
        /* Keep a seperate dictionary, to calculate which coordinates have a set threat level
         * a.k.a total distance to start
         * This allows to find the next coordinate to check fairly quickly, by just
         * looking at the minmum key in this.
         * I initially tried a PriorityQueue, but my implementation was pretty slow and it spent too much time sorting
         */
        var threatLevelDistrubution = [ThreatLevel: Set<Coord>]()
        let startNode = Coord(row: 0, col: 0)
        totalThreatLevel[startNode] = 0
        threatLevelDistrubution[Int.max] = unvisited
        threatLevelDistrubution[Int.max]!.remove(startNode)
        threatLevelDistrubution[0] = Set([startNode])
        
        while let threatLevel = threatLevelDistrubution.keys.min() {
            let currentNode = threatLevelDistrubution[threatLevel]!.first!
            unvisited.remove(currentNode)
            threatLevelDistrubution[threatLevel]!.remove(currentNode)
            
            // Remove key from dictionary if not coords are left for a specific threat level
            if threatLevelDistrubution[threatLevel]!.isEmpty {
                threatLevelDistrubution.removeValue(forKey: threatLevel)
            }
            
            if neighbors[currentNode] == nil {
                neighbors[currentNode] = currentNode
                    .neighbors()
                    .filter {
                        $0.col >= 0 && $0.col <= lastCol && $0.row >= 0 && $0.row <= lastRow
                    }
            }
            var adjacent = neighbors[currentNode]!
            neighbors[currentNode] = adjacent
            adjacent = adjacent
                .filter {
                    unvisited.contains($0)
                }
            let currentThreatLevel = totalThreatLevel[currentNode, default: Int.max]
            for a in adjacent {
                let testCost = parts[a] + currentThreatLevel
                let currentCost = totalThreatLevel[a, default: Int.max]
                if currentCost > testCost {
                    totalThreatLevel[a] = testCost
                    threatLevelDistrubution[currentCost]!.remove(a)
                    threatLevelDistrubution[testCost, default: Set<Coord>()].insert(a)
                    
                    // Remove key from dictionary if not coords are left for a specific threat leveli
                    if threatLevelDistrubution[currentCost]!.isEmpty {
                        threatLevelDistrubution.removeValue(forKey: currentCost)
                    }
                }
            }
        }
        print(totalThreatLevel[Coord(row: lastRow, col: lastCol)]!)
    }
}

extension Array where Element == Array<Int> {
    subscript(coord: Coord) -> Int {
        get { self[coord.row][coord.col] }
        set(newValue) {
            self[coord.row][coord.col] = newValue
        }
    }
}