generalize

solutions/src/2024/16.hs

@@ -56,13 +56,12 @@ Maintainer  : [email protected]
 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]