src/advent06/advent06.hs

   1 import qualified Data.Vector as V
   2 import Data.Vector ((//), (!))
   3 import Data.List (unfoldr)
   4 import qualified Data.Map.Strict as M
   5
   6 type Memory = V.Vector Int
   7 type Redist = (Int, Int, Memory)
   8 type History = M.Map Memory Int
   9
  10 main :: IO ()
  11 main = do
  12         text <- readFile "data/advent06.txt"
  13         let memory = map read $ words text
  14         print $ part1 memory
  15         print $ part2 memory
  16
  17 part1 :: [Int] -> Int
  18 part1 = (findRepeat M.empty) . (zip [0..]) . redistSeq . V.fromList
  19
  20 part2 :: [Int] -> Int
  21 part2 = (findRepeatB M.empty) . (zip [0..]) . redistSeq . V.fromList
  22
  23 findRepeat :: History -> [(Int, Memory)] -> Int
  24 findRepeat h ((n, x) : nxs) = if x `M.member` h
  25                               then n + 1
  26                               else findRepeat (M.insert x n h) nxs
  27
  28 findRepeatB :: History -> [(Int, Memory)] -> Int
  29 findRepeatB h ((n, x) : nxs) = if x `M.member` h
  30                               then n - n0
  31                               else findRepeatB (M.insert x n h) nxs
  32     where n0 = (M.!) h x
  33
  34
  35 redistSeq :: Memory -> [Memory]
  36 redistSeq = unfoldr redistU
  37     where redistU vec = Just (redist vec, redist vec)
  38
  39 redist :: Memory -> Memory
  40 redist = redistR . redistStart
  41
  42
  43
  44 redistStart :: Memory -> Redist
  45 redistStart vec0 = (current, toDistribute, startVec)
  46     where origin = V.maxIndex vec0
  47           toDistribute = vec0!origin
  48           current = (origin + 1) `mod` (length vec0)
  49           startVec = vec0 // [(origin, 0)]
  50
  51 redistR :: Redist -> Memory
  52 redistR (_, 0, vec) = vec
  53 redistR (i, n, vec) = redistR (i', n', vec')
  54     where n' = n - 1
  55           i' = (i + 1) `mod` (length vec)
  56           vec' = vec // [(i, vec!i + 1)]