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