Day 15

1 files changed, 92 insertions, 0 deletions

diff --git a/day15/main.hs b/day15/main.hs
new file mode 100644
index 0000000..cff4233
--- /dev/null
+++ b/day15/main.hs
@@ -0,0 +1,92 @@
+module Main where
+
+import Control.Monad (foldM)
+import Data.Bifunctor (bimap, first, second)
+import Data.Char (digitToInt)
+import Data.Function (on)
+import qualified Data.List as L
+import Data.Map.Strict (Map, (!))
+import qualified Data.Map.Strict as M
+import Data.PQueue.Prio.Min (MinPQueue)
+import qualified Data.PQueue.Prio.Min as PQ
+import Debug.Trace (traceShow, traceShowId)
+import Utils
+
+type Point = (Int, Int)
+
+type Dim = (Int, Int)
+
+type WMap = Map Point Int
+
+type TMap = Map Point Int
+
+update :: WMap -> TMap -> MinPQueue Int Point -> Point -> Point -> (TMap, MinPQueue Int Point)
+update wmap tmap unvisitedPQ source target =
+  let w = wmap ! target
+      oldT = tmap ! target
+      newT = tmap ! source + w
+   in if newT < oldT
+        then (M.insert target (min oldT newT) tmap, PQ.insert newT target unvisitedPQ)
+        else (tmap, unvisitedPQ)
+
+getNeighbors :: WMap -> Point -> [Point]
+getNeighbors wmap p@(x, y) =
+  [(1, 0), (0, 1), (-1, 0), (0, -1)]
+    $> map (bimap (x +) (y +))
+    .> filter (`M.member` wmap)
+
+recurse :: Point -> WMap -> TMap -> MinPQueue Int Point -> Int
+recurse goal wmap tmap unvisitedPQ
+  | null unvisitedPQ = tmap ! goal
+  | otherwise =
+    let ((currentT, current), unvisitedPQ') = PQ.deleteFindMin unvisitedPQ
+
+        -- if currentT == tmap ! current
+        unvisitedNeighbors = getNeighbors wmap current
+        (tmap', unvisitedPQ'') =
+          L.foldl' (\(tm_, pq_) nb -> update wmap tm_ pq_ current nb) (tmap, unvisitedPQ') unvisitedNeighbors
+     in if currentT > tmap ! current
+          then recurse goal wmap tmap unvisitedPQ'
+          else recurse goal wmap tmap' unvisitedPQ''
+
+dijkstra :: Point -> WMap -> Point -> Int
+dijkstra goal costs start =
+  let initTMap = M.mapWithKey (\_ _ -> maxBound :: Int) costs $> M.adjust (const 0) start
+      initUnvisited = PQ.fromList [(0, start)]
+   in recurse goal costs initTMap initUnvisited
+
+e1 chart (xSize, ySize) =
+  let start = (0, 0)
+      goal = (xSize - 1, ySize - 1)
+   in dijkstra goal chart start
+
+e2 chart (xSize, ySize) =
+  let start = (0, 0)
+      goal = (5 * xSize - 1, 5 * ySize - 1)
+      chunks = [(cx, cy) | cx <- [0 .. 4], cy <- [0 .. 4]]
+      chart' =
+        chunks
+          $> foldr
+            ( \(cx, cy) acc ->
+                M.foldrWithKey
+                  ( \(x, y) v ->
+                      M.insert (x + cx * xSize, y + cy * ySize) ((v + cx + cy - 1) `mod` 9 + 1)
+                  )
+                  acc
+                  chart
+            )
+            M.empty
+   in dijkstra goal chart' start
+
+main :: IO ()
+main =
+  do
+    contents <- getContents
+
+    let input = contents $> lines .> map (map digitToInt)
+    let dim@(xSize, ySize) = (length $ head input, length input)
+    let coords = [(x, y) | y <- [0 .. ySize - 1], x <- [0 .. xSize - 1]]
+    let chart = zip coords (concat input) $> M.fromList
+
+    e1 chart dim $> print
+    e2 chart dim $> print