src/advent17/advent17leaky.hs

   1 {-# LANGUAGE NegativeLiterals #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 {-# LANGUAGE OverloadedStrings #-}
   4 {-# LANGUAGE TypeFamilies #-}
   5 {-# LANGUAGE NegativeLiterals #-}
   6 {-# LANGUAGE FlexibleContexts #-}
   7 {-# LANGUAGE OverloadedStrings #-}
   8 {-# LANGUAGE TypeFamilies #-}
   9 {-# LANGUAGE BangPatterns #-}
  10
  11 import Prelude hiding ((++))
  12 import Data.Vector.Unboxed ((!), (++))
  13 import qualified Data.Vector.Unboxed as V
  14 import qualified Data.List as L
  15
  16 type RingBuffer = V.Vector Int
  17
  18 initialStepSize :: Int
  19 initialStepSize = 366
  20
  21 initialRingBuffer :: RingBuffer
  22 initialRingBuffer = V.singleton 0
  23
  24
  25 main :: IO ()
  26 main = do
  27         print $ part1 initialStepSize 2017
  28         print $ part2 initialStepSize 50000000
  29
  30
  31 part1 n k = (!1) $! last $! take k $! L.unfoldr ringBufferUnfolder (1, n, initialRingBuffer)
  32
  33
  34 part2 n k = finalBuffer!targetLoc
  35     where finalBuffer = last $! take k $! L.unfoldr ringBufferUnfolder (1, n, initialRingBuffer)
  36           zeroLoc = V.head $ V.elemIndices 0 finalBuffer
  37           targetLoc = (zeroLoc + 1) `rem` (V.length finalBuffer)
  38
  39
  40 updateRingBuffer nextValue stepSize buffer = (buffer', stepSize)
  41     where !nextPos = (stepSize + 1) `rem` V.length buffer
  42           (!start, !end) = V.splitAt nextPos buffer
  43           !buffer' = V.cons nextValue $ end ++ start
  44
  45 ringBufferUnfolder :: (Int, Int, RingBuffer) -> Maybe (RingBuffer, (Int, Int, RingBuffer))
  46 ringBufferUnfolder (nextValue, stepSize, buffer) = Just (buffer', (nextValue + 1, stepSize', buffer'))
  47     where (!buffer', !stepSize') = updateRingBuffer nextValue stepSize buffer