From 62e6b042370e600819a29609c3b9e9f09326a71f Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Wed, 15 Dec 2021 21:23:57 +0100
Subject: 15

---
 2021/15.hs | 103 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 103 insertions(+)
 create mode 100644 2021/15.hs

(limited to '2021/15.hs')

diff --git a/2021/15.hs b/2021/15.hs
new file mode 100644
index 0000000..542ba42
--- /dev/null
+++ b/2021/15.hs
@@ -0,0 +1,103 @@
+{-# LANGUAGE TupleSections #-}
+module Main where
+
+import Control.Monad (forM_)
+import Control.Monad.ST (ST)
+import qualified Data.Array as A
+import qualified Data.Array.ST as MA
+import Data.Function (on)
+import Data.List hiding (insert)
+import Data.Ord
+
+-- import Debug.Trace
+
+import Input
+
+
+-- Priority queue that only performs okay if there are only a small number of
+-- distinct keys in the queue at any one time.
+newtype PQueue k v = PQueue [(k, [v])]
+  deriving (Show)
+qFromList :: Ord k => [(k, v)] -> PQueue k v
+qFromList [] = PQueue []
+qFromList l = PQueue (map ((,) <$> fst . head <*> map snd)
+                      . groupBy ((==) `on` fst)
+                      $ sortBy (comparing fst) l)
+qAdds :: Ord k => [(k, v)] -> PQueue k v -> PQueue k v
+qAdds = \pairs (PQueue l) -> PQueue (inserts (let PQueue pq = qFromList pairs in pq) l)
+  where inserts :: Ord k => [(k, [v])] -> [(k, [v])] -> [(k, [v])]
+        inserts [] l = l
+        inserts news [] = news
+        inserts news@((k, vs) : restnews) l@((k', vs') : rest) = case compare k k' of
+          LT -> (k, vs) : inserts restnews l
+          EQ -> (k, vs ++ vs') : inserts restnews rest
+          GT -> (k', vs') : inserts news rest
+qAdd :: Ord k => k -> v -> PQueue k v -> PQueue k v
+qAdd k v pq = qAdds [(k, v)] pq
+qPop :: PQueue k v -> Maybe ((k, v), PQueue k v)
+qPop (PQueue []) = Nothing
+qPop (PQueue ((k, [v]) : l)) = Just ((k, v), PQueue l)
+qPop (PQueue ((k, v:vs) : l)) = Just ((k, v), PQueue ((k, vs) : l))
+qPop (PQueue ((_, []) : l)) = qPop (PQueue l)
+
+-- Return array: (parent position, minimal distance)
+dijkstra :: A.Array (Int, Int) Int -> A.Array (Int, Int) ((Int, Int), Int)
+dijkstra board = MA.runSTArray $ do
+    arr <- MA.newArray (A.bounds board) ((-1, -1), maxBound)
+    MA.writeArray arr (0, 0) ((-1, -1), 0)
+    loop arr (qFromList [(0, (0, 0))])
+    return arr
+  where
+    loop :: MA.STArray s (Int, Int) ((Int, Int), Int) -> PQueue Int (Int, Int) -> ST s ()
+    loop arr qu
+      | Just ((curdist, (x, y)), qu') <- qPop qu = do
+          -- traceM ("pq = " ++ show qu)
+          -- traceM ("visit " ++ show (x, y))
+          -- traceM ("pq = " ++ show qu')
+          neighbours <- mapM (\pos -> (pos,) . snd <$> MA.readArray arr pos)
+                        . filter (A.inRange (A.bounds board))
+                        $ [(x+1,y), (x,y+1), (x-1,y), (x,y-1)]
+          let updates = [(curdist + board A.! pos', pos')
+                        | (pos', dist) <- neighbours
+                        , dist > curdist + board A.! pos']
+          -- traceM ("updates = " ++ show updates)
+          forM_ updates $ \(newdist, pos') ->
+              MA.writeArray arr pos' ((x, y), newdist)
+          -- MA.freeze arr >>= traceM . printDists
+          loop arr (qAdds updates qu')
+      | otherwise = return ()
+
+printDists :: A.Array (Int, Int) ((Int, Int), Int) -> String
+printDists dists =
+    let pad2 s = let s' = "  " ++ s in drop (length s' - 2) s'
+        (_, (ymax, xmax)) = A.bounds dists
+    in unlines [intercalate " " [let d = snd (dists A.! (y, x))
+                                 in if d == maxBound then "--" else pad2 (show d)
+                                | x <- [0..xmax]]
+               | y <- [0..ymax]]
+
+tracePath :: A.Array (Int, Int) ((Int, Int), Int) -> (Int, Int) -> [((Int, Int), Int)]
+tracePath _ (-1, -1) = []
+tracePath dists to = let (next, dist) = dists A.! to
+                     in (to, dist) : tracePath dists next
+
+expand :: Int -> (Int -> a -> a) -> A.Array (Int, Int) a -> A.Array (Int, Int) a
+expand times f arr =
+    let ((0, 0), (ymax, xmax)) = A.bounds arr
+        (h, w) = (ymax + 1, xmax + 1)
+    in A.listArray ((0, 0), (times * h - 1, times * w - 1))
+         [f (x `div` w + y `div` w) (arr A.! (y `mod` h, x `mod` w))
+         | y <- [0 .. times * h - 1]
+         , x <- [0 .. times * w - 1]]
+
+main :: IO ()
+main = do
+    board <- map (map (read . pure)) <$> getInput 15
+    let (h, w) = (length board, length (head board))
+    let boardarr = A.listArray ((0, 0), (h-1, w-1)) (concat board)
+    print (snd (dijkstra boardarr A.! (h-1, w-1)))
+    let expanded = expand 5 (\n x -> (x + n - 1) `mod` 9 + 1) boardarr
+    -- putStr $ unlines [concatMap show [expanded A.! (y, x) | x <- [0 .. 5*w-1]] | y <- [0 .. 5*h-1]]
+    print (snd (dijkstra expanded A.! (5*h-1, 5*w-1)))
+    -- mapM_ print $ reverse (tracePath dists (h-1, w-1))
+    -- putStr (printDists dists)
-- 
cgit v1.2.3-70-g09d2