advent-2021/9/solution.hs

71 lines
2.3 KiB
Haskell

import Data.List
import Data.List.Split
import Data.Maybe
main :: IO ()
main = do
input <- readFile "9/input.txt"
let x = parse $ lines input
-- Problem 1
print $ f x
-- Problem 2
print $ g x
where
f = risk
g = solve
parse :: [String] -> [[Int]]
parse = map (map read . chunksOf 1)
risk :: [[Int]] -> Int
risk a = sum $ map succ $ mapMaybe index (areas a)
where
index = ($ a) . uncurry idx
solve :: [[Int]] -> Int
solve x = product $ take 3 $ reverse $ sort $ map (basin x) $ areas x
-- | Find risk areas (low points).
areas :: [[Int]] -> [(Int,Int)]
areas a = [ (x,y) | y <- [ 0 .. (pred $ length a) ],
x <- [ 0 .. (length (a !! y)) ],
Just n <- [idx x y a],
all ((> n) . snd) (adj x y a) ]
-- | Get the size of a basin centered on the given point.
basin :: [[Int]] -> (Int, Int) -> Int
basin m (x,y) = length a
where
a = case idx x y m of
Just v -> let t = ((x,y),v) in t : search t m
Nothing -> []
-- | Search for neighbors of the given node where the value of the neighbor
-- | is not 9 and is greater than the value of the given node. This search
-- | is then repeated for each result it finds until no nodes satisfying the
-- | constraint can be found in the immediate neighborhood of the cluster.
search :: ((Int, Int), Int) -> [[Int]] -> [((Int, Int), Int)]
search t @ ((x,y),v) b = nub (next ++ (next >>= (`search` b)))
where
neighbors = adj x y b
next = nub $ filter (\(_,n) -> n /= 9 && n > v) neighbors
-- | Find elements that are directly adjacent to a point in a matrix.
adj :: Int -> Int -> [[a]] -> [((Int, Int), a)]
adj x y m = [ ((a, b), v) | a <- [pred x .. succ x],
b <- [pred y .. succ y],
x /= a || y /= b,
y == b || x == a,
Just v <- [idx a b m] ]
-- | Index into a matrix safely. If out of bounds on either axis,
-- | returns `Nothing`.
idx :: Int -> Int -> [[a]] -> Maybe a
idx x y m = get y m >>= get x
-- | Index into a list safely. Returns `Nothing` if out of bounds.
get :: Int -> [a] -> Maybe a
get i l | i < 0 = Nothing
| i >= length l = Nothing
| otherwise = Just (l !! i)