diff options
Diffstat (limited to '2021/15.hs')
| -rw-r--r-- | 2021/15.hs | 103 |
1 files changed, 103 insertions, 0 deletions
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) |