78 lines
2.0 KiB
Haskell
78 lines
2.0 KiB
Haskell
import Data.Char (digitToInt)
|
|
import Data.List (nub, sortBy)
|
|
import Data.Ord
|
|
|
|
main :: IO ()
|
|
main = do
|
|
input <- getContents
|
|
let
|
|
heightMap = parseHeights input
|
|
(putStrLn . show . solution1) heightMap
|
|
(putStrLn . show . solution2) heightMap
|
|
|
|
solution1 :: [[Int]] -> Int
|
|
solution1 heights = sum (map (\p -> heights !!! p + 1) (lowPoints heights))
|
|
|
|
solution2 :: [[Int]] -> Int
|
|
solution2 heights = product (take 3 (sortBy (comparing Down) (map length basins)))
|
|
where
|
|
basins :: [[(Int, Int)]]
|
|
basins = map (flip basin heights) (lowPoints heights)
|
|
|
|
basin :: (Int, Int) -> [[Int]] -> [(Int, Int)]
|
|
basin point heights
|
|
| null next = []
|
|
| otherwise = nub (point:(next)++flatmap (flip basin heights) next)
|
|
where
|
|
next = filter partOfBasin (adjacent point heights)
|
|
partOfBasin x =
|
|
let
|
|
xHeight = heights !!! x
|
|
in
|
|
xHeight > heights !!! point && xHeight /= 9
|
|
|
|
flatmap :: (t -> [a]) -> [t] -> [a]
|
|
flatmap _ [] = []
|
|
flatmap f (x:xs) = f x ++ flatmap f xs
|
|
|
|
lowPoints :: [[Int]] -> [(Int, Int)]
|
|
lowPoints heights = filter (flip isLowPoint heights) (indices heights)
|
|
where
|
|
indices :: [[Int]] -> [(Int, Int)]
|
|
indices m = [(x, y) |
|
|
x <- [0..length(m)-1],
|
|
y <- [0..length(m !! 0)-1]]
|
|
|
|
isLowPoint :: (Int, Int) -> [[Int]] -> Bool
|
|
isLowPoint point m = all (> (m !!! point)) (map ((!!!) m) (adjacent point m))
|
|
|
|
(!!!) :: [[Int]] -> (Int, Int) -> Int
|
|
(!!!) m index = m !! (fst index) !! (snd index)
|
|
|
|
adjacent :: (Int, Int) -> [[Int]] -> [(Int, Int)]
|
|
adjacent (x, y) m =
|
|
[(i+x, j+y) |
|
|
i <- [-1..1],
|
|
j <- [-1..1],
|
|
abs i /= abs j,
|
|
i+x >= 0 && i+x < length(m),
|
|
j+y >= 0 && j+y < length(m !! 0)]
|
|
|
|
parseHeights :: String -> [[Int]]
|
|
parseHeights = map (map digitToInt) . lines
|
|
|
|
-- Tests
|
|
|
|
testInput = unlines [
|
|
"2199943210",
|
|
"3987894921",
|
|
"9856789892",
|
|
"8767896789",
|
|
"9899965678"
|
|
]
|
|
|
|
testInputParsed = parseHeights testInput
|
|
|
|
test1 = solution1 testInputParsed -- 15
|
|
test2 = solution2 testInputParsed -- 1134
|