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
104
105
106
107
108
109
|
-- requires cabal install --lib megaparsec parser-combinators heap vector
module Main (main) where
import Control.Applicative.Permutations
import Control.Monad (void, when)
import qualified Data.Char as C
import Data.Either
import Data.Functor
import qualified Data.Heap as H
import qualified Data.List as L
import qualified Data.Map as M
import Data.Maybe
import qualified Data.Set as S
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import Data.Void (Void)
import Text.Megaparsec
import Text.Megaparsec.Char
import Debug.Trace
exampleExpectedOutput = 94
type Line = VU.Vector Int
type Input = V.Vector Line
type Parser = Parsec Void String
parseLine :: Parser Line
parseLine = do
line <- some (C.digitToInt <$> digitChar) <* eol
return $ VU.generate (length line) (line !!)
parseInput' :: Parser Input
parseInput' = do
line <- some parseLine <* eof
return $ V.generate (length line) (line !!)
parseInput :: String -> IO Input
parseInput filename = do
input <- readFile filename
case runParser parseInput' filename input of
Left bundle -> error $ errorBundlePretty bundle
Right input' -> return input'
data Heading = NS | EW deriving (Eq, Show)
data Position = Position Int Int Int Heading deriving (Show) -- cost x y heading
instance Ord Position where
compare (Position c1 _ _ _) (Position c2 _ _ _) = c1 `compare` c2
instance Eq Position where
(Position c1 _ _ _) == (Position c2 _ _ _) = c1 == c2
type Candidates = H.MinHeap Position
type CostsState = M.Map (Int,Int) (Int, Int) -- (x, y) (NSCost, EWCost)
step :: Input -> CostsState -> Position -> [Position]
step input costs (Position cost x y h) = L.concat [next 1 1 0, next (-1) (-1) 0]
where
next :: Int -> Int -> Int -> [Position] -- diff increment costchange
next 11 _ _ = []
next (-11) _ _ = []
next d i costChange = case heatLoss of
Just hl -> if improvement hl then (Position (cost + costChange + hl) x' y' h') : next (d+i) i (costChange + hl)
else next (d+i) i (costChange + hl)
Nothing -> []
where
h' | h == NS = EW
| otherwise = NS
x' | h' == NS = x
| otherwise = x + d
y' | h' == NS = y + d
| otherwise = y
improvement :: Int -> Bool
improvement hl | d < 4 && d > -4 = False
| otherwise = case M.lookup (x', y') costs of
Just (h1, h2) -> case h' of
NS -> cost + hl < h1
EW -> cost + hl < h2
Nothing -> undefined -- should not happen, we catch out of bound when looking for heatLoss
heatLoss :: Maybe Int
heatLoss = case input V.!? y' of
Just line -> line VU.!? x'
Nothing -> Nothing
compute :: Input -> Int
compute input = compute' startingCosts startingCandidates
where
compute' :: CostsState -> Candidates -> Int
compute' costs candidates | x == size - 1 && y == size - 1 = cost
| improvement = compute' costs' (H.union candidates' (H.fromList $ step input costs position))
| otherwise = compute' costs candidates'
where
candidates' = H.drop 1 candidates
costs' = M.insert (x, y) (if h == NS then (cost, snd state) else (fst state, cost)) costs
improvement | h == NS = cost < fst state
| otherwise = cost < snd state
position@(Position cost x y h) = head $ H.take 1 candidates
state = costs M.! (x, y)
infinity = maxBound :: Int
startingCandidates = H.fromList [Position 0 0 0 h|h <- [NS, EW]] :: Candidates
startingCosts = M.fromList [((x, y), (infinity, infinity))|x<-[0..size], y<-[0..size]]
size = V.length input
main :: IO ()
main = do
example <- parseInput "example"
let exampleOutput = compute example
when (exampleOutput /= exampleExpectedOutput) (error $ "example failed: got " ++ show exampleOutput ++ " instead of " ++ show exampleExpectedOutput)
input <- parseInput "input"
print $ compute input
|
