diff options
| -rw-r--r-- | 10-2.hs | 27 |
1 files changed, 23 insertions, 4 deletions
@@ -1,8 +1,11 @@ +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 !)) @@ -19,9 +22,19 @@ width,height :: [[a]] -> Int width = length . head height = length -score :: Map -> Pos -> Int -score m p | m ! p == 9 = 1 - | otherwise = sum . map (score m) $ step m p +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 @@ -34,4 +47,10 @@ parse ls = listArray main :: IO () main = do m <- parse . lines <$> getContents - print . sum . map (score m) . trailheads $ m + print + . sum + . toList + . flip evalState (listArray (bounds m) $ repeat Nothing) + . traverse (score' m) + . trailheads + $ m |