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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
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)
|
