Day sixteen part two

day-16-haskell/src/Lib.hs

@@ -46,41 +46,64 @@ task1 input =
   let worldMap = parseFile input
       rawGraph = mapToGraph worldMap
   in do
-    --rawGraph |> wEdges |> Map.elems |> join |> showEdgesSimple |> putStrLn
     worldGraph <- simplifyGraph rawGraph
     worldGraph |> wEdges |> Map.elems |> join |> showEdgesSimple |> putStrLn
-    --return $ expectJust "No path found" $ findBestPath worldGraph [(wStart worldGraph, East)]
-    bestPath <- dijkstra worldGraph Map.empty (Map.singleton (wStart worldGraph, East) 0)
-    return $ expectJust "No path found" $ bestPath
+    (bestPath, _) <- dijkstra worldGraph Map.empty (Map.singleton (wStart worldGraph, East) (0, []))
+    return bestPath
 
 task2 :: String -> IO IntType
-task2 _ = return 2
+task2 input = 
+  let worldMap = parseFile input
+      rawGraph = mapToGraph worldMap
+  in do
+    (bestCost, bestPaths) <- dijkstra rawGraph Map.empty (Map.singleton (wStart rawGraph, East) (0, []))
+    [wEnd rawGraph] : bestPaths |> join |> Set.fromList |> length |> return
 
 showEdgesSimple :: [Edge] -> String
 showEdgesSimple [] = ""
 showEdgesSimple (edge:edges) = (show $ eFrom edge) ++ " -> " ++ (show $ eTo edge) ++ ": " ++ (show $ ePrice edge) ++ "\n" ++ (showEdgesSimple edges)
 
-dijkstra :: WorldGraph -> Map.Map (Vec2, Direction) Int -> Map.Map (Vec2, Direction) Int -> IO (Maybe Int)
-dijkstra world@WorldGraph { wEdges, wEnd } done unvisited
+type WasBestePreis = (Int, [(Vec2, Direction)])
+dijkstra :: WorldGraph -> Map.Map (Vec2, Direction) WasBestePreis -> Map.Map (Vec2, Direction) WasBestePreis -> IO (Int, [[Vec2]])
+dijkstra world@WorldGraph { wEdges, wEnd, wStart } done unvisited
   | null unvisited =
-    [North, East, South, West] |> map (\d -> done Map.!? (wEnd, d) |> maybeToList) |> join |> minimum |> pure |> return
+    let bestePreis = [North, East, South, West] |> map (\d -> done Map.!? (wEnd, d) |> maybeToList) |> join |> foldr combineBestePreis (maxBound :: IntType, [])
+    in bestePreis |> \(p, os) -> (p, map (reconstructPath done) os |> join) |> return
   | otherwise =
-    let ((pos, dir), cost) = findNext
+    let ((pos, dir), (cost, _)) = findNext
+        nextDone = Map.insert (pos, dir) (unvisited Map.! (pos, dir)) done
         outgoing = Map.findWithDefault [] pos wEdges
         unvisitedOutgoing = outgoing |> filter (\Edge { eTo, eDirTo } -> Map.notMember (eTo, eDirTo) done)
         unvisitedOutgoingWithPrice = unvisitedOutgoing |> map (\e@Edge { eDirFrom, ePrice } -> (e, if dir == dOpposite eDirFrom then cost + ePrice else cost + ePrice + 1000)) :: [(Edge, IntType)]
         unvisitedWithoutCurrent = Map.delete (pos, dir) unvisited
-        nextUnvisited = foldr (\(Edge { eTo, eDirTo }, price) -> Map.alter (pure . maybe price (min price)) (eTo, eDirTo)) unvisitedWithoutCurrent unvisitedOutgoingWithPrice
-        nextDone = Map.insert (pos, dir) cost done
-    in do
-      -- putStrLn $ "Processing " ++ (show pos) ++ "/" ++ (show dir)
-      -- print nextUnvisited
-      dijkstra world nextDone nextUnvisited
+        nextUnvisited = foldr (\(Edge { eTo, eDirTo }, price) -> Map.alter (pure . maybeCombineBestePreis (price, [(pos, dir)])) (eTo, eDirTo)) unvisitedWithoutCurrent unvisitedOutgoingWithPrice
+    in dijkstra world nextDone nextUnvisited
   where
-    findNext :: ((Vec2, Direction), Int)
+    findNext :: ((Vec2, Direction), WasBestePreis)
     findNext =
       let options = unvisited |> Map.assocs
-      in foldr (\a@(_,aC) b@(_,bC) -> if aC < bC then a else b) (head options) (drop 1 options)
+      in foldr (\a@(_,(aC,_)) b@(_,(bC,_)) -> if aC < bC then a else b) (head options) (drop 1 options)
+
+    maybeCombineBestePreis :: WasBestePreis -> Maybe WasBestePreis -> WasBestePreis
+    maybeCombineBestePreis a Nothing = a
+    maybeCombineBestePreis a (Just b) = combineBestePreis a b
+
+    combineBestePreis :: WasBestePreis -> WasBestePreis -> WasBestePreis
+    combineBestePreis (aP, aOs) (bP, bOs)
+      | aP < bP = (aP, aOs)
+      | aP == bP = (aP, aOs ++ bOs)
+      | aP > bP = (bP, bOs)
+
+reconstructPath :: Map.Map (Vec2, Direction) WasBestePreis -> (Vec2, Direction) -> [[Vec2]]
+reconstructPath done cur@(curPos,_) =
+  case done Map.!? cur of
+    Nothing -> [[curPos]]
+    Just (_,[]) -> [[curPos]]
+    Just (_,os) -> map (reconstructPath done) os |> join |> map (\subPath -> curPos : subPath)
+
+showEntries :: Show a => Show b => [(a,b)] -> String
+showEntries [] = ""
+showEntries ((k,v):rest) = show k ++ ": " ++ show v ++ "\n" ++ showEntries rest
 
 findBestPath :: WorldGraph -> [(Vec2, Direction)] -> Maybe Int
 findBestPath _ [] = error "Must call findBestPath with an initial position"