1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
|
{-# 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 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
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)
dijkstra :: A.Array (Int, Int) Int -> A.Array (Int, Int) Int
dijkstra board = MA.runSTArray $ do
arr <- MA.newArray (A.bounds board) maxBound
MA.writeArray arr (0, 0) 0
loop arr (qFromList [(0, (0, 0))])
return arr
where
loop :: MA.STArray s (Int, Int) Int -> PQueue Int (Int, Int) -> ST s ()
loop arr qu
| Just ((curdist, (x, y)), qu') <- qPop qu = do
neighbours <- mapM (\pos -> (pos,) <$> 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']
forM_ updates $ \(newdist, pos') ->
MA.writeArray arr pos' newdist
loop arr (qAdds updates qu')
| otherwise = return ()
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 (dijkstra boardarr A.! (h-1, w-1))
let expanded = expand 5 (\n x -> (x + n - 1) `mod` 9 + 1) boardarr
print (dijkstra expanded A.! (5*h-1, 5*w-1))
|
