Main.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
module Main where

import           Control.Lens
import           Control.Monad.State

import qualified Data.IntMap.Strict  as IM
import           Data.List.Split     (splitOn)

main :: IO ()
main = do
  nums <- fmap read . splitOn "," <$> readFile "input.txt"
  print (ageMap nums 80)
  print (ageMap nums 256)

ageMap :: [Int] -> Int -> Int
ageMap xs n = IM.foldl' (+) 0 $ execState (replicateM_ n step) initial
  where
    initial = IM.fromListWith (+) (zip xs (repeat 1))

step :: State (IM.IntMap Int) ()
step = mapM_ (uncurry next) =<< IM.assocs <$> get

next :: Int -> Int -> State (IM.IntMap Int) ()
next 0 v = do
  at 6 . non 0 += v
  at 0 . non 0 -= v
  at 8 . non 0 += v
next k v = do
  at (k - 1) . non 0 += v
  at k . non 0 -= v