advent23/src/advent23.hs

   1 -- import Debug.Trace
   2
   3 import qualified Data.List.PointedList.Circular as P
   4 import Control.Lens
   5 import Data.Maybe
   6 import Data.Char (digitToInt)
   7
   8 import Control.Monad
   9 import Control.Monad.ST
  10 import Data.STRef
  11 import qualified Data.Vector.Unboxed.Mutable as V
  12
  13
  14 puzzleInput :: [Int]
  15 puzzleInput = map digitToInt "538914762"
  16 -- puzzleInput = map digitToInt "389125467" -- example
  17
  18 main :: IO ()
  19 main =
  20   do  putStrLn $ part1 puzzleInput
  21       print $ part2 puzzleInput
  22
  23
  24 part1 nums = label $ playN cups 100
  25   where cups = fromJust $ P.fromList nums
  26
  27 part2 nums = a * b
  28   where finalCups = runGame nums 1000000 10000000
  29         (a, b) = clockwiseOf1 finalCups
  30
  31
  32 label cups = concatMap show $ tail $ takeRight (P.length cups) atOne
  33   where atOne = fromJust $ P.find 1 cups
  34
  35 playN cups n = (iterate playOne cups) !! n
  36
  37 playOne cups = P.next replacedAtCurrent
  38   where current = cups ^. P.focus
  39         held = takeRight 3 $ P.next cups
  40         shorter = fromJust $ dropRight 3 $ P.next cups
  41         destination = validDestination (current - 1) 9 held
  42         shorterAtDestination = fromJust $ P.find destination shorter
  43         replaced = foldr P.insertRight shorterAtDestination $ reverse held
  44         replacedAtCurrent = fromJust $ P.find current replaced
  45
  46
  47 validDestination 0 max missing = validDestination max max missing
  48 validDestination n max missing
  49   | n `elem` missing = validDestination (n - 1) max missing
  50   | otherwise = n
  51
  52
  53 takeRight :: Int -> P.PointedList a -> [a]
  54 takeRight 0 _ = []
  55 takeRight n xs = (xs ^. P.focus):(takeRight (n - 1) $ P.next xs)
  56
  57 dropRight :: Int -> P.PointedList a -> Maybe (P.PointedList a)
  58 dropRight 0 xs = Just xs
  59 dropRight n xxs = case (P.deleteRight xxs) of
  60   Just xs -> dropRight (n - 1) xs
  61   Nothing -> Nothing
  62
  63
  64 clockwiseOf1 cups = (a, b)
  65   where a = cups!!1
  66         b = cups!!a
  67
  68 runGame :: [Int] -> Int -> Int -> [Int]
  69 runGame seed cupsNeeded roundsNeeded =
  70   runST $
  71     do cups <- seedGame seed cupsNeeded
  72        gameLoop roundsNeeded cupsNeeded cups
  73        mapM (V.read cups) [0..cupsNeeded]
  74
  75 seedGame :: [Int] -> Int -> ST s (V.MVector s Int)
  76 seedGame seed cupsNeeded =
  77   do  cups <- V.new (cupsNeeded + 1)
  78       let extended = seed ++ [10, 11..]
  79       forM_ [0..cupsNeeded] $ \i -> V.write cups i (i + 1)
  80       forM_ (zip seed $ tail extended) $ \(i, j) -> V.write cups i j
  81       V.write cups 0 (head seed)
  82       let end = if cupsNeeded > (length seed)
  83                 then cupsNeeded
  84                 else last seed
  85       V.write cups end (head seed)
  86       return cups
  87
  88 gameLoop targetRound maxCups cups =
  89     do forM_ [1..targetRound]
  90              (\_ -> gameStep maxCups cups)
  91        return ()
  92
  93 gameStep :: Int -> V.MVector s Int -> ST s ()
  94 gameStep maxCups cups =
  95   do  current <- V.read cups 0
  96       held1 <- V.read cups current
  97       held2 <- V.read cups held1
  98       held3 <- V.read cups held2
  99       afterHeld <- V.read cups held3
 100
 101           -- close the loop, removing the held cups
 102       V.write cups current afterHeld
 103
 104       let destination =
 105             validDestination (current - 1) maxCups [held1, held2, held3]
 106       afterDestination <- V.read cups destination
 107
 108           -- make the held come after the destination
 109       V.write cups destination held1
 110
 111           -- make the end of the held point into the rest of the loop
 112       V.write cups held3 afterDestination
 113
 114           -- advance current
 115       nextCup <- V.read cups current
 116           -- and store it
 117       V.write cups 0 nextCup
 118       return ()
 119