import qualified Data.Array as A import Data.Array (Array, (!)) import Data.Foldable (toList) import qualified Data.Set as Set import Data.Bifunctor (bimap) amap :: A.Ix i => (i -> a -> b) -> Array i a -> Array i b amap f arr = A.listArray (A.bounds arr) (zipWith f (A.indices arr) (toList arr)) main :: IO () main = do topo' <- map (map (read @Int . pure)) . lines <$> getContents let w = length (head topo') let h = length topo' let topo = A.listArray ((0, 0), (w - 1, h - 1)) (concat topo') let flood9 = amap (\i n -> if n == 9 then (Set.singleton i, 1::Int) else (mempty, 0)) topo let inBounds (x, y) = 0 <= x && x < w && 0 <= y && y < h let neighs (x, y) = filter inBounds [(x-1,y), (x+1,y), (x,y-1), (x,y+1)] let stepdown mp = amap (\i n -> bimap Set.unions sum $ unzip [mp ! p | p <- neighs i, topo ! p == n + 1]) topo let flood = iterate stepdown flood9 !! 9 let (totscore, totrank) = bimap sum sum $ unzip $ map (\((s,n),_) -> (Set.size s, n)) $ filter (\(_,n) -> n == 0) $ zip (A.elems flood) (A.elems topo) print totscore print totrank