summaryrefslogtreecommitdiff
path: root/10-2.hs
blob: 6ee5bc25dff579850a097b31b379ade2300604ee (plain)
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
import Control.Monad.State
import Data.Array
import Data.Foldable (toList)

data Pos = Pos Int Int deriving (Eq,Ix,Ord,Show)

type Map = Array Pos Int
type S = State (Array Pos (Maybe Int)) Int

step :: Map -> Pos -> [Pos]
step m p@(Pos x y) = filter ((== 1) . subtract (m ! p) . (m !))
  . filter (`elem` indices m)
  $ steps
  where
    steps = [ Pos (x+1) y
            , Pos (x-1) y
            , Pos x (y+1)
            , Pos x (y-1)
            ]

width,height :: [[a]] -> Int
width = length . head
height = length

score :: Map -> Pos -> S
score m p | m ! p == 9 = return 1
          | otherwise  = fmap sum . traverse (score' m) $ step m p

score' :: Map -> Pos -> S
score' m p = do
  memo <- get
  case memo ! p of
    Just x  -> return x
    Nothing -> do
      x <- score m p
      modify (// [(p,Just x)])
      return x

trailheads :: Map -> [Pos]
trailheads = map fst . filter ((== 0) . snd) . assocs

parse :: [[Char]] -> Map
parse ls = listArray
           (Pos 0 0, Pos (width ls - 1) (height ls - 1))
           (map (read . pure) $ concat ls)

main :: IO ()
main = do
  m <- parse . lines <$> getContents
  print
    . sum
    . toList
    . flip evalState (listArray (bounds m) $ repeat Nothing)
    . traverse (score' m)
    . trailheads
    $ m