advent06/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/06/advent-of-code-2021-day-6/
   2
   3 import Data.List
   4 import Data.List.Split
   5 import qualified Data.IntMap.Strict as M
   6
   7 type Shoal = M.IntMap Integer
   8
   9 main :: IO ()
  10 main =
  11   do  text <- readFile "data/advent06.txt"
  12       let fish = map read $ splitOn "," text
  13       let shoal = mkShoal fish
  14       let generations = iterate generation shoal
  15       print $ part1 generations
  16       print $ part2 generations
  17
  18 part1 :: [Shoal] -> Integer
  19 part1 = countInGeneration 80
  20
  21 part2 :: [Shoal] -> Integer
  22 part2 = countInGeneration 256
  23
  24 countInGeneration :: Int -> [Shoal] -> Integer
  25 countInGeneration n generations = sum $ M.elems (generations!!n)
  26
  27 generation :: Shoal -> Shoal
  28 generation shoal = M.union (age shoal) (birth shoal)
  29
  30 age :: Shoal -> Shoal
  31 age shoal = M.insert 6 (was0 + was7) shoal'
  32   where shoal' = M.mapKeys ageFish shoal
  33         was0 = M.findWithDefault 0 0 shoal
  34         was7 = M.findWithDefault 0 7 shoal
  35
  36 ageFish :: Int -> Int
  37 ageFish 0 = 6
  38 ageFish n = n - 1
  39
  40 birth :: Shoal -> Shoal
  41 birth shoal = M.singleton 8 parents
  42   where parents = M.findWithDefault 0 0 shoal
  43
  44 mkShoal :: [Int] -> Shoal
  45 mkShoal = M.fromList
  46   . map (\g -> (head g, fromIntegral $ length g))
  47   . group
  48   . sort