track predecessors

solutions/src/2024/16.hs

@@ -64,6 +64,7 @@ import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Set (Set)
 import Data.Set qualified as Set
+import Advent.Search (fillN)
 
 -- | >>> :main
 -- 88416
@@ -74,50 +75,54 @@ main =
     let start:_ = [p | (p,'S') <- assocs input]
         end  :_ = [p | (p,'E') <- assocs input]
         open    = amap ('#' /=) input
-        
+
         -- start with all the possible initial facings so that the optimization later
         -- that assumes we only turn 90-degrees before moving will hold
-        path0 = Set.singleton start
-        q0 = IntMap.fromList [(   0, Map.singleton (start, east) path0)
-                             ,(1000, Map.fromList [((start, north), path0), ((start, south), path0)])
-                             ,(2000, Map.singleton (start, west) path0)]
+        q0 = IntMap.fromList [(   0, Map.singleton (start, east) [])
+                             ,(1000, Map.fromList [((start, north), []), ((start, south), [])])
+                             ,(2000, Map.singleton (start, west) [])]
+
+        (cost, preds) = shortestPath open end Map.empty q0
+        shortestPathNodes
+          = Set.map fst
+          $ fillN (preds Map.!)
+            [(end,v) | v <- [north, east, south, west], Map.member (end, v) preds]
 
-        (p1, p2) = search open end Set.empty q0
-    print p1
-    print p2
+    print cost
+    print (length shortestPathNodes)
 
-search ::
-  UArray Coord Bool                       {- ^ grid marking open spaces -} ->
-  Coord                                   {- ^ target coordinate -} ->
-  Set (Coord, Coord)                      {- ^ position/velocity pairs already finished -} ->
-  IntMap (Map (Coord, Coord) (Set Coord)) {- ^ cost to (position/velocity to nodes-on-path -} ->
-  (Int, Int)                              {- ^ cost of shortest path and nodes on shorts paths -}
-search open end seen q =
+-- | Mapping from a node to the predecessors of that node on the shortest path to that node
+type Predecessors = Map (Coord, Coord) [(Coord, Coord)]
+
+-- | 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)
-      | null dones -> search open end seen' q2
-      | otherwise  -> (cost, Set.size (Set.unions dones))
+      | 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.withoutKeys states seen
+        states' = Map.difference states seen
 
         -- look for states that have reached the target
-        dones = [visited | ((p, _), visited) <- Map.assocs states', p == end]
+        done = any (\(p,_) -> p == end) (Map.keys states')
 
         -- mark all the new states at this cost as seen so we don't revisit them again
-        seen' = Set.union seen (Map.keysSet states')
+        seen' = Map.union seen states'
 
         -- queue up all the successor states to be visited in the future
-        q2 = IntMap.unionWith merge q1
-           $ IntMap.fromListWith merge
-              [ next
-                | ((p, v), path) <- Map.assocs states'
-                , next <- [(cost + 1001, Map.singleton (p', v') (Set.insert p' path)) | let v' = turnRight v, let p' = p + v', open ! p']
-                       ++ [(cost + 1001, Map.singleton (p', v') (Set.insert p' path)) | let v' = turnLeft  v, let p' = p + v', open ! p']
-                       ++ [(cost +    1, Map.singleton (p', v ) (Set.insert p' path)) |                       let p' = p + v , open ! p']
-              ]
-
-        merge = Map.unionWith Set.union
-
-
+        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']
+          ]