{-|
Module      : Main
Description : Day 9 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/9>

Find the basins on the height map.

>>> :{
exampleGrid =
  Data.Array.Unboxed.listArray
    (C 0 0, C 4 9)
    [2,1,9,9,9,4,3,2,1,0,
     3,9,8,7,8,9,4,9,2,1,
     9,8,5,6,7,8,9,8,9,2,
     8,7,6,7,8,9,6,7,8,9,
     9,8,9,9,9,6,5,6,7,8]
:}

>>> exampleBasins = toBasinIds exampleGrid

>>> basinRiskSum exampleGrid exampleBasins
15

>>> basinSizes exampleBasins
[3,9,14,9]

-}
module Main (main) where

import Advent (counts, getInputArray, arrIx)
import Advent.Coord (Coord(..), cardinal)
import Data.Array.Unboxed (Array, UArray, (!), amap, array, assocs, bounds, elems)
import Data.Char (digitToInt)
import Data.Foldable (toList)
import Data.List (sortBy)
import Data.Maybe (catMaybes)

-- | >>> :main
-- 588
-- 964712
main :: IO ()
IO ()
main =
 do UArray Coord Int
heights <- (Char -> Int) -> UArray Coord Char -> UArray Coord Int
forall (a :: * -> * -> *) e' e i.
(IArray a e', IArray a e, Ix i) =>
(e' -> e) -> a i e' -> a i e
amap Char -> Int
digitToInt (UArray Coord Char -> UArray Coord Int)
-> IO (UArray Coord Char) -> IO (UArray Coord Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2021 Int
9
    let basinIds :: BasinIds
basinIds = UArray Coord Int -> BasinIds
toBasinIds UArray Coord Int
heights
    Int -> IO ()
forall a. Show a => a -> IO ()
print (UArray Coord Int -> BasinIds -> Int
basinRiskSum UArray Coord Int
heights BasinIds
basinIds)
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Int -> [Int] -> [Int]
forall a. Ord a => Int -> [a] -> [a]
top Int
3 (BasinIds -> [Int]
basinSizes BasinIds
basinIds)))

-- * Basins

-- | Array of locations in the grid and the unique low point they flow towards.
type BasinIds = Array Coord (Maybe Coord)

-- | Compute the basin low points for each location on the grid.
toBasinIds :: UArray Coord Int -> BasinIds
toBasinIds :: UArray Coord Int -> BasinIds
toBasinIds UArray Coord Int
heights = BasinIds
basinIds
  where
    basinIds :: BasinIds
basinIds = (Coord, Coord) -> [(Coord, Maybe Coord)] -> BasinIds
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (UArray Coord Int -> (Coord, Coord)
forall i. Ix i => UArray i Int -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Coord Int
heights) [(Coord
c, Coord -> Int -> Maybe Coord
basinId Coord
c Int
h) | (Coord
c, Int
h) <- UArray Coord Int -> [(Coord, Int)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Coord Int
heights]
    
    basinId :: Coord -> Int -> Maybe Coord
basinId Coord
_ Int
9 = Maybe Coord
forall a. Maybe a
Nothing -- problem defines height 9 not to be in a basin
    basinId Coord
c Int
h =
     do [Coord]
xs <- [Maybe Coord] -> Maybe [Coord]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [BasinIds
basinIdsBasinIds -> Coord -> Maybe Coord
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Coord
x | Coord
x <- Coord -> [Coord]
cardinal Coord
c, Int
xh <- UArray Coord Int -> Coord -> [Int]
forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx UArray Coord Int
heights Coord
x, Int
xhInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
h]
        case [Coord]
xs of
          []                  -> Coord -> Maybe Coord
forall a. a -> Maybe a
Just Coord
c -- this is the lowest point in the basin
          Coord
y:[Coord]
ys | (Coord -> Bool) -> [Coord] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Coord
yCoord -> Coord -> Bool
forall a. Eq a => a -> a -> Bool
==) [Coord]
ys -> Coord -> Maybe Coord
forall a. a -> Maybe a
Just Coord
y -- all flows go to a unique basin
               | Bool
otherwise    -> Maybe Coord
forall a. Maybe a
Nothing -- no unique basin

-- | Compute the sum of the risk values of the basin low points.
basinRiskSum :: UArray Coord Int -> BasinIds -> Int
basinRiskSum :: UArray Coord Int -> BasinIds -> Int
basinRiskSum UArray Coord Int
heights BasinIds
basinIds = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
h | (Coord
c, Int
h) <- UArray Coord Int -> [(Coord, Int)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Coord Int
heights, Coord -> Maybe Coord
forall a. a -> Maybe a
Just Coord
c Maybe Coord -> Maybe Coord -> Bool
forall a. Eq a => a -> a -> Bool
== BasinIds
basinIdsBasinIds -> Coord -> Maybe Coord
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Coord
c]

-- | List of the sizes of all the basins on the grid.
basinSizes :: BasinIds -> [Int]
basinSizes :: BasinIds -> [Int]
basinSizes = Map Coord Int -> [Int]
forall a. Map Coord a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Map Coord Int -> [Int])
-> (BasinIds -> Map Coord Int) -> BasinIds -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Coord] -> Map Coord Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts ([Coord] -> Map Coord Int)
-> (BasinIds -> [Coord]) -> BasinIds -> Map Coord Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe Coord] -> [Coord]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Coord] -> [Coord])
-> (BasinIds -> [Maybe Coord]) -> BasinIds -> [Coord]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasinIds -> [Maybe Coord]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
elems

-- * List utilities

-- | Returns the @n@ largest elements of a list
top :: Ord a => Int -> [a] -> [a]
top :: forall a. Ord a => Int -> [a] -> [a]
top Int
n = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((a -> a -> Ordering) -> a -> a -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare)