advent21/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/26/advent-of-code-2021-day-21/
   2
   3 import Debug.Trace
   4
   5 import Data.Text ()
   6 import qualified Data.Text.IO as TIO
   7 import Data.Attoparsec.Text hiding (take, takeWhile)
   8 import Control.Applicative
   9
  10 import qualified Data.Map.Strict as M
  11 import Data.Map.Strict ((!))
  12 import Data.List
  13 import qualified Data.MultiSet as MS
  14
  15 data Player = Player1 | Player2 deriving (Eq, Ord, Show)
  16
  17 data PlayerState = PlayerState
  18   { position :: Int
  19   , score :: Int
  20   } deriving (Eq, Ord, Show)
  21
  22 data Game = Game
  23   { players :: M.Map Player PlayerState
  24   , current :: Player
  25   , rolls :: Int
  26   } deriving (Eq, Ord)
  27
  28 type Games = MS.MultiSet Game
  29 type Dice = MS.MultiSet Int
  30 type Winners = MS.MultiSet Player
  31
  32 instance Show Game where
  33   show game = "{" ++ (showPlayer Player1) ++ (showActive) ++ (showPlayer Player2) ++ "}"
  34     where showPlayer p = (show $ position $ (players game) ! p) ++ "," ++ (show $ score $ (players game) ! p)
  35           showActive = if (current game) == Player1 then "<" else ">"
  36
  37
  38 main :: IO ()
  39 main =
  40   do  text <- TIO.readFile "data/advent21.txt"
  41       let game = successfulParse text
  42       print $ part1 game
  43       print $ part2 game
  44
  45 part1 game = scoreGame finalGame
  46   where finalGame = head $ dropWhile (not . finished 1000) $ scanl' gameStep game detDice
  47         detDice = map (\n -> sum ([d `mod1` 100 | d <- [n..(n+2)]]::[Int])) [1, 4..]
  48
  49 part2 game = max (Player1 `MS.occur` winners) (Player2 `MS.occur` winners)
  50   where games0 = MS.singleton game
  51         winners0 = MS.empty
  52         winners = nonDetGameSimulation 21 games0 diracDice winners0
  53
  54 finished :: Int -> Game -> Bool
  55 finished threshold game = any (>= threshold) $ map score $ M.elems (players game)
  56
  57 scoreGame :: Game -> Int
  58 scoreGame game = (rolls game) * losingScore
  59   where losingScore = minimum $ map score $ M.elems (players game)
  60
  61 diracDice :: Dice
  62 diracDice = MS.fromList [a + b + c | a <- [1..3], b <- [1..3], c <- [1..3]]
  63
  64 gameStep :: Game -> Int -> Game
  65 gameStep game theseRolls = game'
  66   where activePlayer = (players game) ! (current game)
  67         pos = position activePlayer
  68         sc = score activePlayer
  69         pos' = (pos + theseRolls) `mod1` 10
  70         sc' = sc + pos'
  71         activePlayer' = PlayerState {position = pos', score = sc'}
  72         current' = nextPlayer (current game)
  73         players' = M.insert (current game) activePlayer' (players game)
  74         game' = Game { players = players'
  75                      , current = current'
  76                      , rolls = rolls game + 3
  77                      }
  78
  79 nonDetGameSimulation :: Int -> Games -> Dice -> Winners -> Winners
  80 nonDetGameSimulation winThreshold games0 dice winners0
  81   -- | trace ((show games0) ++ "; " ++ (show winners0)) False = undefined
  82   -- | trace (show winners0) False = undefined
  83   | MS.null games0 = winners0
  84   | otherwise = nonDetGameSimulation winThreshold games dice winners
  85   where games' = nonDetGameStep games0 dice
  86         (winGames, games) = MS.partition (finished winThreshold) games'
  87         p1Wins = MS.size $ MS.filter (\g -> current g == Player2) winGames
  88         p2Wins = MS.size $ MS.filter (\g -> current g == Player1) winGames
  89         winners = MS.insertMany Player2 p2Wins $ MS.insertMany Player1 p1Wins winners0
  90
  91 nonDetGameStep :: Games -> Dice -> Games
  92 -- nonDetGameStep games dice | trace ("G0 >" ++ (show games) ++ "-" ++ (show dice)) False = undefined
  93 nonDetGameStep games dice = MS.foldOccur (nonDetGameStep1 dice) MS.empty games
  94
  95 nonDetGameStep1 :: Dice -> Game -> MS.Occur -> Games -> Games
  96 -- nonDetGameStep1 dice game gnum acc | trace ("G1 >" ++ (show game) ++ "-" ++ (show dice) ++ ": " ++ (show gnum)) False = undefined
  97 nonDetGameStep1 dice game gnum acc = MS.foldOccur (nonDetGameStep2 game gnum) acc dice
  98
  99 nonDetGameStep2 :: Game -> MS.Occur -> Int -> MS.Occur -> Games -> Games
 100 -- nonDetGameStep2 dice dnum game gnum acc | trace ("G2 >" ++ (show game) ++ "-" ++ (show dice) ++ ": " ++ (show gnum) ++ "," ++ (show dnum)) False = undefined
 101 nonDetGameStep2 game gnum roll dnum acc = MS.insertMany game' (gnum * dnum) acc
 102   where game' = gameStep game roll
 103
 104
 105 nextPlayer :: Player -> Player
 106 nextPlayer Player1 = Player2
 107 nextPlayer Player2 = Player1
 108
 109 mod1 :: Int -> Int -> Int
 110 mod1 a b = ((a - 1) `mod` b) + 1
 111
 112
 113
 114 -- Parsing
 115
 116 playerP = ("1" *> pure Player1) <|> ("2" *> pure Player2)
 117
 118 playerStateP = playerify <$> ("Player " *> playerP) <*> (" starting position: " *> decimal)
 119   where playerify name pos = (name, PlayerState {position = pos, score = 0})
 120
 121 gameP = gamify <$> playerStateP `sepBy` endOfLine
 122   where gamify ps = Game {rolls = 0, current = Player1, players = M.fromList ps}
 123
 124
 125 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 126 successfulParse input =
 127   case parseOnly gameP input of
 128     Left  _err ->  Game {rolls=0, current=Player1, players=M.empty} -- TIO.putStr $ T.pack $ parseErrorPretty err
 129     Right game -> game
 130