21.hs

{-# Language BlockArguments, ImportQualifiedPost, QuasiQuotes, GeneralisedNewtypeDeriving #-}
{-# OPTIONS_GHC -w #-}
{-|
Module      : Main
Description : Day 21 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/21>

-}
module Main (main) where

import Advent (counts, format)
import Advent.Memo (memo4)
import Control.Applicative (Alternative((<|>)))
import Control.Monad (replicateM)
import Control.Monad.Trans.Writer.CPS (runWriterT, writerT, WriterT)
import Data.Coerce (coerce)
import Data.Map (Map)
import Data.Map.Strict qualified as Map
import Data.Monoid (Product(Product))
import Data.List (unfoldr, partition)
import Control.Monad (guard)

-- | >>> :main
-- 428736
-- 57328067654557
main :: IO ()
main =
 do (p1,p2) <- [format|21 Player 1 starting position: %u%nPlayer 2 starting position: %u%n|]
    print (part1 0 p1 p2 0 0)
    print (part2 p1 p2)
    -- print (maximum (snd <$> runPaths (part2 p1 p2 0 0)))

-- | Compute the @die rolls * losing score@ once one player
-- wins with 1000 points.
part1 ::
  Int {- ^ turn counter -} ->
  Int {- ^ player 1 location -} ->
  Int {- ^ player 2 location -} ->
  Int {- ^ player 1 score    -} ->
  Int {- ^ player 2 score    -} ->
  Int {- ^ player 2 score * 3 * turns -}
part1 turns p1 p2 p1s p2s
  | p1s' >= 1000 = 3 * turns' * p2s
  | otherwise    = part1 turns' p2 p1' p2s p1s'
  where
    turns' = turns + 1
    p1'    = wrap (p1 + 6 - turns) 10
    p1s'   = p1s + p1'

part2 ::
  Int {- player 1's starting location -} ->
  Int {- player 2's starting location -} ->
  Int {- ways player 1 can win -}
part2 p1 p2 = sum (zipWith (*) p1Wins (1 : p2Live))
  where
    p1Wins = unfoldr p2step (Map.singleton (p1,0) 1)
    p2Wins = unfoldr p2step (Map.singleton (p2,0) 1)
    p2Live = scanl (\acc w -> acc * 27 - w) 1 (tail p2Wins)

p2step :: Map.Map (Int, Int) Int -> Maybe (Int, Map.Map (Int, Int) Int)
p2step games
  | Map.null games = Nothing
  | otherwise = Just (sum wins, live')
  where
    (wins, live) = Map.partitionWithKey (\(_,score) _ -> score >= 21) games
    live' = Map.fromListWith (+)
              do ((loc, score),n1) <- Map.toList live
                 (roll,n2) <- Map.toList threeRolls
                 let loc' = wrap (loc+roll) 10
                 pure ((loc', score + loc'), n1*n2)

-- | Sum of 3d3.
threeRolls :: Map Int Int
threeRolls = counts (sum <$> replicateM 3 [1..3])

-- * Modular arithmetic

-- | Wrap number between @1@ and an inclusive upper bound
wrap :: Int {- ^ value -} -> Int {- ^ bound -} -> Int
wrap x n = (x - 1) `mod` n + 1