93 lines
3.6 KiB
Haskell
93 lines
3.6 KiB
Haskell
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
|
|
|
|
module Day18.Puzzle2 (day18_2) where
|
|
|
|
import qualified Data.Array as A
|
|
import qualified Data.HashMap.Strict as M
|
|
import Data.List.Split (splitOn)
|
|
import Data.Maybe (fromJust)
|
|
import qualified Data.PSQueue as PQ
|
|
|
|
type Coords = (Int, Int)
|
|
|
|
newtype Graph = Graph {edges :: M.HashMap Coords [Coords]} deriving (Show)
|
|
|
|
data Distance a = Dist a | Infinity deriving (Eq)
|
|
|
|
instance (Ord a) => Ord (Distance a) where
|
|
Infinity <= Infinity = True
|
|
Infinity <= Dist _ = False
|
|
Dist _ <= Infinity = True
|
|
Dist x <= Dist y = x <= y
|
|
|
|
instance (Show a) => Show (Distance a) where
|
|
show Infinity = "Infinity"
|
|
show (Dist x) = show x
|
|
|
|
addDistance :: (Num a) => Distance a -> Distance a -> Distance a
|
|
addDistance (Dist x) (Dist y) = Dist (x + y)
|
|
addDistance _ _ = Infinity
|
|
|
|
data DijkstraState = DijkstraState
|
|
{ unvisited :: PQ.PSQ Coords (Distance Int),
|
|
distances :: M.HashMap Coords (Distance Int)
|
|
}
|
|
deriving (Show)
|
|
|
|
updateDistances :: M.HashMap Coords (Distance Int) -> [Coords] -> Distance Int -> M.HashMap Coords (Distance Int)
|
|
updateDistances dists [] _ = dists
|
|
updateDistances dists (n : nodes) startD =
|
|
updateDistances (M.adjust (const startD) n dists) nodes startD
|
|
|
|
visit :: PQ.PSQ Coords (Distance Int) -> Coords -> [Coords] -> Distance Int -> PQ.PSQ Coords (Distance Int)
|
|
visit us node [] _ = PQ.delete node us
|
|
visit us node (e : es) dist = visit (PQ.adjust (const dist) e us) node es dist
|
|
|
|
visitNode :: DijkstraState -> Graph -> Coords -> Distance Int -> DijkstraState
|
|
visitNode state graph node d =
|
|
let es = edges graph M.! node
|
|
ds = updateDistances (distances state) es d
|
|
us = visit (unvisited state) node es d
|
|
in state {unvisited = us, distances = ds}
|
|
|
|
findShortestPath :: Graph -> Coords -> Coords -> Distance Int
|
|
findShortestPath graph start end =
|
|
let nodesDist = (start PQ.:-> Dist 0) : [k PQ.:-> Infinity | k <- M.keys $ edges graph, k /= start]
|
|
dists = (start, Dist 0) : [(k, Infinity) | k <- M.keys $ edges graph, k /= start]
|
|
initialState = DijkstraState {unvisited = PQ.fromList nodesDist, distances = M.fromList dists}
|
|
in dijkstra initialState
|
|
where
|
|
dijkstra s =
|
|
let nd = fromJust $ PQ.findMin (unvisited s)
|
|
n = PQ.key nd
|
|
d = PQ.prio nd
|
|
in if n == end
|
|
then d
|
|
else
|
|
if d == Infinity
|
|
then Infinity
|
|
else dijkstra $ visitNode s graph n (addDistance d (Dist 1))
|
|
|
|
adjacent :: A.Array Coords Char -> Coords -> Coords -> [Coords]
|
|
adjacent array (i, j) (maxI, maxJ) = [(a, b) | (a, b) <- [(i, j + 1), (i, j - 1), (i + 1, j), (i - 1, j)], a >= 0, b >= 0, a <= maxI, b <= maxJ, array A.! (a, b) /= '#']
|
|
|
|
corruptMemory :: A.Array Coords Char -> [Coords] -> A.Array Coords Char
|
|
corruptMemory = foldl (\a b -> a A.// [(b, '#')])
|
|
|
|
findFirstBlocker :: A.Array Coords Char -> [Coords] -> Coords -> Coords -> Coords
|
|
findFirstBlocker memory (c:cs) start end =
|
|
let memory' = corruptMemory memory [c]
|
|
memoryGraph = Graph {edges = M.fromList [(k, adjacent memory' k (70, 70)) | k <- A.indices memory']}
|
|
in if findShortestPath memoryGraph start end == Infinity
|
|
then c
|
|
else findFirstBlocker memory' cs start end
|
|
|
|
day18_2 :: IO ()
|
|
day18_2 = do
|
|
contents <- map (splitOn ",") . lines <$> readFile "input/day18.txt"
|
|
let memory = A.listArray ((0, 0), (70, 70)) $ replicate 5041 '.'
|
|
coords = take 1024 [(read x, read y) | (x : y : _) <- contents]
|
|
coords' = drop 1024 [(read x, read y) | (x : y : _) <- contents]
|
|
memory' = corruptMemory memory coords
|
|
putStrLn $ "Day 18, Puzzle 2 solution: " ++ show (findFirstBlocker memory' coords' (0, 0) (70, 70))
|