54 lines
2.1 KiB
Haskell
54 lines
2.1 KiB
Haskell
module Day09 where
|
|
import Data.Char ( digitToInt )
|
|
import Data.Maybe (catMaybes)
|
|
import Data.Containers (IsMap(filterMap))
|
|
import qualified Data.HashSet as HashSet
|
|
import Data.HashSet (HashSet)
|
|
import Data.List (sort)
|
|
|
|
infixr 5 !?
|
|
(!?) :: [a] -> Int -> Maybe a
|
|
(!?) [] _ = Nothing
|
|
(!?) (x:xs) i
|
|
| i < 0 = Nothing
|
|
| i == 0 = Just x
|
|
| otherwise = xs !? i-1
|
|
|
|
getGrid :: [[Int]] -> Int -> Int -> Maybe Int
|
|
getGrid grid x y = do
|
|
row <- grid !? y
|
|
row !? x
|
|
|
|
getMinimumNeighbor :: [[Int]] -> Int -> Int -> Int
|
|
getMinimumNeighbor grid x y = minimum $ catMaybes [g (x-1) y, g (x+1) y, g x (y-1), g x (y+1)]
|
|
where g x y = getGrid grid x y
|
|
|
|
part1:: [[Int]] -> Int
|
|
part1 input = sum $ map sum $ zipWith (\y row -> zipWith (\x digit -> apply x y digit) [0 .. ] row) [0..] input
|
|
where apply x y digit = if digit < getMinimumNeighbor input x y then digit + 1 else 0
|
|
|
|
getLows :: [[Int]] -> [(Int,Int)]
|
|
getLows input = concatMap catMaybes (zipWith (\y row -> zipWith (\x digit -> apply x y digit) [0 .. ] row) [0..] input)
|
|
where apply x y digit = if digit < getMinimumNeighbor input x y then Just (x,y) else Nothing
|
|
|
|
stepBasin :: [[Int]] -> HashSet (Int,Int) -> HashSet (Int,Int)
|
|
stepBasin input set = HashSet.foldr HashSet.union HashSet.empty $ HashSet.map (\(x,y) -> HashSet.fromList $ getElems x y) set
|
|
where
|
|
getElems x y = (x,y) : map (\(x,y,n) -> (x,y)) (getHigherNeighs (Just $ input !! y !! x) x y)
|
|
getHigherNeighs digit x y = filter (\(x,y,n) -> n > digit && n < Just 9) $ getNeighs x y
|
|
getNeighs x y = [g (x-1) y, g (x+1) y, g x (y-1), g x (y+1)]
|
|
g x y = (x,y,getGrid input x y)
|
|
|
|
getBasin :: [[Int]] -> (Int,Int) -> Int
|
|
getBasin input (x,y) = length $ HashSet.toList $ calcBasin $ HashSet.fromList [(x,y)]
|
|
where
|
|
calcBasin set = let nextSet = stepBasin input set in if length nextSet == length set then set else calcBasin nextSet
|
|
|
|
part2 :: [[Int]] -> Int
|
|
part2 input = product $ take 3 $ reverse $ sort $ map (getBasin input) $ getLows input
|
|
|
|
main :: IO ()
|
|
main = do
|
|
input <- map (map digitToInt) . lines <$> getContents
|
|
putStrLn $ "part 1: " ++ show (part1 input)
|
|
putStrLn $ "part 2: " ++ show (part2 input) |