diff --git a/solutions/src/2024/16.hs b/solutions/src/2024/16.hs index 8e62ca9..1c39fa2 100644 --- a/solutions/src/2024/16.hs +++ b/solutions/src/2024/16.hs @@ -56,13 +56,12 @@ Maintainer : emertens@gmail.com module Main (main) where import Advent (getInputArray) -import Advent.Coord (Coord, north, east, south, west, turnRight, turnLeft) -import Data.Array.Unboxed (UArray, assocs, amap, (!)) +import Advent.Coord (north, east, south, west, turnRight, turnLeft) +import Data.Array.Unboxed (assocs, amap, (!)) import Data.IntMap (IntMap) import Data.IntMap qualified as IntMap import Data.Map (Map) import Data.Map qualified as Map -import Data.Set (Set) import Data.Set qualified as Set import Advent.Search (fillN) @@ -82,47 +81,48 @@ main = ,(1000, Map.fromList [((start, north), []), ((start, south), [])]) ,(2000, Map.singleton (start, west) [])] - (cost, preds) = shortestPath open end Map.empty q0 - shortestPathNodes + step (p,v) = [(1001, (p', v')) | let v' = turnRight v, let p' = p + v', open ! p'] + ++ [(1001, (p', v')) | let v' = turnLeft v, let p' = p + v', open ! p'] + ++ [( 1, (p', v )) | let p' = p + v , open ! p'] + isEnd (p, _) = p == end + + (cost, preds) = shortestPath step isEnd q0 + nodesOnShortestPaths = Set.map fst $ fillN (preds Map.!) [(end,v) | v <- [north, east, south, west], Map.member (end, v) preds] print cost - print (length shortestPathNodes) + print (length nodesOnShortestPaths) -- | Mapping from a node to the predecessors of that node on the shortest path to that node -type Predecessors = Map (Coord, Coord) [(Coord, Coord)] +type Predecessors a = Map a [a] -- | Main loop for a shortest-path implementation that computes the cost of the shortest path shortestPath :: - UArray Coord Bool {- ^ grid marking open spaces -} -> - Coord {- ^ target coordinate -} -> - Predecessors {- ^ finalized optimal predecessors -} -> - IntMap Predecessors {- ^ cost to candidate predecessor additions -} -> - (Int, Predecessors) {- ^ cost of shortest path to end and predecessors of all shortest paths -} -shortestPath open end seen q = - case IntMap.minViewWithKey q of - Nothing -> error "no solution" - Just ((cost, states), q1) - | done -> (cost, seen') - | otherwise -> shortestPath open end seen' q2 - where - -- remove all the states at this cost that we've seen at a lower cost - states' = Map.difference states seen + Ord a => + (a -> [(Int, a)]) {- ^ successors of a node -} -> + (a -> Bool) {- ^ predicate for the destination -} -> + IntMap (Predecessors a) {- ^ cost to candidate predecessor additions -} -> + (Int, Predecessors a) {- ^ cost of shortest path to end and predecessors of all shortest paths -} +shortestPath = go Map.empty + where + go seen step isEnd q = + case IntMap.minViewWithKey q of + Nothing -> error "no solution" + Just ((cost, states), q1) + | done -> (cost, seen') + | otherwise -> go seen' step isEnd q2 + where + -- remove all the states at this cost that we've seen at a lower cost + states' = Map.difference states seen - -- look for states that have reached the target - done = any (\(p,_) -> p == end) (Map.keys states') + -- look for states that have reached the target + done = any isEnd (Map.keys states') - -- mark all the new states at this cost as seen so we don't revisit them again - seen' = Map.union seen states' + -- mark all the new states at this cost as seen so we don't revisit them again + seen' = Map.union seen states' - -- queue up all the successor states to be visited in the future - q2 = foldl (\m (k,v) -> IntMap.insertWith (Map.unionWith (++)) k v m) q1 - [ (cost + amt, Map.singleton next [cur]) - | cur@(p, v) <- Map.keys states' - , (amt, next) - <- [(1001, (p', v')) | let v' = turnRight v, let p' = p + v', open ! p'] - ++ [(1001, (p', v')) | let v' = turnLeft v, let p' = p + v', open ! p'] - ++ [( 1, (p', v )) | let p' = p + v , open ! p'] - ] + -- queue up all the successor states to be visited in the future + q2 = foldl (\m (k,v) -> IntMap.insertWith (Map.unionWith (++)) k v m) q1 + [(cost + amt, Map.singleton next [cur]) | cur <- Map.keys states', (amt, next) <- step cur]