-- compile with -O2 to take advantage of memoization! import Data.Array import Data.List (nub) data Pos = Pos Int Int deriving (Eq,Ix,Ord,Show) type Map = Array Pos Int type Memo = Array Pos [Pos] memo :: Map -> Memo memo m = array (bounds m) [ (p, reachable m p) | p <- indices m ] 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 reachable :: Map -> Pos -> [Pos] reachable m p | m ! p == 9 = [p] | otherwise = nub . concatMap (memo m !) $ step m p 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 . map (length . (memo m !)) . trailheads $ m