1 -- Writeup at https://work.njae.me.uk/2021/12/06/advent-of-code-2021-day-6/
5 import qualified Data.IntMap.Strict as M
7 type Shoal = M.IntMap Integer
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
18 part1 :: [Shoal] -> Integer
19 part1 = countInGeneration 80
21 part2 :: [Shoal] -> Integer
22 part2 = countInGeneration 256
24 countInGeneration :: Int -> [Shoal] -> Integer
25 countInGeneration n generations = sum $ M.elems (generations!!n)
27 generation :: Shoal -> Shoal
28 generation shoal = M.union (age shoal) (birth 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
40 birth :: Shoal -> Shoal
41 birth shoal = M.singleton 8 parents
42 where parents = M.findWithDefault 0 0 shoal
44 mkShoal :: [Int] -> Shoal
46 . map (\g -> (head g, fromIntegral $ length g))