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
|
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
|
