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

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

Play a game on a simple game of play players on a circular board with
9 spaces. The players roll dice to advance to a numbered space and
earn that many points.

Part 1 simulates the game directly and part 2 takes advantage of
there being no interaction between the two players to simulate
their play separately.

-}
module Main (main) where

import Advent (counts, format)
import Control.Applicative (Alternative)
import Control.Monad (replicateM)
import Control.Monad.Trans.Writer.Strict (WriterT(..))
import Data.Coerce (coerce)
import Data.List (scanl', unfoldr)
import Data.Map (Map)
import Data.Map.Strict qualified as Map
import Data.Monoid (Product(..))

-- | >>> :main
-- 428736
-- 57328067654557
main :: IO ()
IO ()
main =
 do (Int
p1,Int
p2) <- [format|2021 21 Player 1 starting position: %u%nPlayer 2 starting position: %u%n|]
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> Int -> Int
part1 Int
p1 Int
p2)
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> Int -> Int
part2 Int
p1 Int
p2)

-- * Part 1

-- | Compute the @die rolls * losing score@ once one player
-- wins with 1000 points.
--
-- >>> part1 4 8
-- 739785
part1 ::
  Int {- ^ player 1 location -} ->
  Int {- ^ player 2 location -} ->
  Int {- ^ player 2 score * roll count -}
part1 :: Int -> Int -> Int
part1 = Int -> Int -> Int -> Int -> Int -> Int
go Int
0 Int
0 Int
0
  where
    go :: Int -> Int -> Int -> Int -> Int -> Int
go Int
turns Int
s1 Int
s2 Int
p1 Int
p2
      | Int
s1' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1000 = Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
turns' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s2
      | Bool
otherwise   = Int -> Int -> Int -> Int -> Int -> Int
go Int
turns' Int
s2 Int
s1' Int
p2 Int
p1'
      where
        turns' :: Int
turns' = Int
turns Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        p1' :: Int
p1'    = Int -> Int -> Int
wrap (Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
6 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
turns) Int
10
        s1' :: Int
s1'    = Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p1'

-- * Part 2

-- | Count the maximum number of ways a player can win in the most given
-- that the game is played rolling 3d3.
--
-- >>> part2 4 8
-- 444356092776315
part2 ::
  Int {- ^ player 1's starting location -} ->
  Int {- ^ player 2's starting location -} ->
  Int {- ^ ways player 1 can win -}
part2 :: Int -> Int -> Int
part2 Int
p1 Int
p2 = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
u1 Int
u2
  where
    w1 :: [Int]
w1 = Int -> [Int]
wins Int
p1                        -- ways player 1 wins by turn
    w2 :: [Int]
w2 = Int -> [Int]
wins Int
p2                        -- ways player 2 wins by turn
    l1 :: [Int]
l1 = [Int] -> [Int]
toLoses [Int]
w1                     -- ways player 1 hasn't won by turn
    l2 :: [Int]
l2 = [Int] -> [Int]
toLoses [Int]
w2                     -- ways player 2 hasn't won by turn
    u1 :: Int
u1 = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
w1 [Int]
l2)        -- universes in which player 1 wins
    u2 :: Int
u2 = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
w2 ([Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
tail [Int]
l1)) -- universes in which player 2 wins

-- | Compute the ways a player can win in part 2 per turn given a starting position.
--
-- >>> wins 4
-- [0,0,4608,249542,3219454,24905476,77993473,62172638,8678745,53217]
wins :: Int -> [Int]
wins :: Int -> [Int]
wins Int
x = (Map (Int, Int) Int -> Maybe (Int, Map (Int, Int) Int))
-> Map (Int, Int) Int -> [Int]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr Map (Int, Int) Int -> Maybe (Int, Map (Int, Int) Int)
p2step ((Int, Int) -> Int -> Map (Int, Int) Int
forall k a. k -> a -> Map k a
Map.singleton (Int
x, Int
0) Int
1)

-- | Compute the ways a player can not win in part 2 per turn given a starting position.
--
-- >>> toLoses (wins 8)
-- [1,27,729,17953,254050,1411009,3520415,2121762,219716,1206,0]
toLoses :: [Int] -> [Int]
toLoses :: [Int] -> [Int]
toLoses = (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' (\Int
acc Int
w -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Map Int Int -> Int
forall a. Num a => Map Int a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Map Int Int
rolls Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w) Int
1

-- | Advance the counts of states by playing one turn of the game.
p2step ::
  Map (Int, Int) Int {- ^ live games ((location, score), ways) -} ->
  Maybe (Int, Map (Int, Int) Int) {- ^ wins and next turn's live games -}
p2step :: Map (Int, Int) Int -> Maybe (Int, Map (Int, Int) Int)
p2step Map (Int, Int) Int
games
  | Map (Int, Int) Int -> Bool
forall k a. Map k a -> Bool
Map.null Map (Int, Int) Int
games = Maybe (Int, Map (Int, Int) Int)
forall a. Maybe a
Nothing
  | Bool
otherwise      = (Int, Map (Int, Int) Int) -> Maybe (Int, Map (Int, Int) Int)
forall a. a -> Maybe a
Just (Map (Int, Int) Int -> Int
forall a. Num a => Map (Int, Int) a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Map (Int, Int) Int
winStates, Map (Int, Int) Int
games')
  where
    (Map (Int, Int) Int
winStates, Map (Int, Int) Int
games') =
      ((Int, Int) -> Int -> Bool)
-> Map (Int, Int) Int -> (Map (Int, Int) Int, Map (Int, Int) Int)
forall k a. (k -> a -> Bool) -> Map k a -> (Map k a, Map k a)
Map.partitionWithKey (\(Int
_,Int
score) Int
_ -> Int
score Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
21) (Map (Int, Int) Int -> (Map (Int, Int) Int, Map (Int, Int) Int))
-> Map (Int, Int) Int -> (Map (Int, Int) Int, Map (Int, Int) Int)
forall a b. (a -> b) -> a -> b
$
      Counter (Int, Int) -> Map (Int, Int) Int
forall a. Ord a => Counter a -> Map a Int
fromCounter
        [ (Int
loc', Int
score Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
loc')
          | (Int
loc, Int
score) <- Map (Int, Int) Int -> Counter (Int, Int)
forall a. Map a Int -> Counter a
toCounter Map (Int, Int) Int
games
          , Int
roll         <- Map Int Int -> Counter Int
forall a. Map a Int -> Counter a
toCounter Map Int Int
rolls
          , let loc' :: Int
loc' = Int -> Int -> Int
wrap (Int
locInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
roll) Int
10
          ]

-- | Outcomes of rolling 3d3
--
-- >>> rolls
-- fromList [(3,1),(4,3),(5,6),(6,7),(7,6),(8,3),(9,1)]
rolls :: Map Int Int
rolls :: Map Int Int
rolls  = [Int] -> Map Int Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [[Int]] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Int] -> [[Int]]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
3 [Int
1..Int
3])

-- * Modular arithmetic

-- | Wrap number between @1@ and an inclusive upper bound.
--
-- >>> [wrap i 4 | i <- [-2..6]]
-- [2,3,4,1,2,3,4,1,2]
wrap :: Int {- ^ value -} -> Int {- ^ bound -} -> Int
wrap :: Int -> Int -> Int
wrap Int
x Int
n = (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- * Tracking counts

-- | Type for backtracking computations that can keep track of how many
-- ways a state is reachable. This allows us to alternate between two
-- useful representations:
--
-- * @[(a,Int)]@ - good for backtracking
-- * @Map a Int@ - good for consolidation
newtype Counter a = Counter (WriterT (Product Int) [] a)
  deriving ((forall a b. (a -> b) -> Counter a -> Counter b)
-> (forall a b. a -> Counter b -> Counter a) -> Functor Counter
forall a b. a -> Counter b -> Counter a
forall a b. (a -> b) -> Counter a -> Counter b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Counter a -> Counter b
fmap :: forall a b. (a -> b) -> Counter a -> Counter b
$c<$ :: forall a b. a -> Counter b -> Counter a
<$ :: forall a b. a -> Counter b -> Counter a
Functor, Functor Counter
Functor Counter =>
(forall a. a -> Counter a)
-> (forall a b. Counter (a -> b) -> Counter a -> Counter b)
-> (forall a b c.
    (a -> b -> c) -> Counter a -> Counter b -> Counter c)
-> (forall a b. Counter a -> Counter b -> Counter b)
-> (forall a b. Counter a -> Counter b -> Counter a)
-> Applicative Counter
forall a. a -> Counter a
forall a b. Counter a -> Counter b -> Counter a
forall a b. Counter a -> Counter b -> Counter b
forall a b. Counter (a -> b) -> Counter a -> Counter b
forall a b c. (a -> b -> c) -> Counter a -> Counter b -> Counter c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> Counter a
pure :: forall a. a -> Counter a
$c<*> :: forall a b. Counter (a -> b) -> Counter a -> Counter b
<*> :: forall a b. Counter (a -> b) -> Counter a -> Counter b
$cliftA2 :: forall a b c. (a -> b -> c) -> Counter a -> Counter b -> Counter c
liftA2 :: forall a b c. (a -> b -> c) -> Counter a -> Counter b -> Counter c
$c*> :: forall a b. Counter a -> Counter b -> Counter b
*> :: forall a b. Counter a -> Counter b -> Counter b
$c<* :: forall a b. Counter a -> Counter b -> Counter a
<* :: forall a b. Counter a -> Counter b -> Counter a
Applicative, Applicative Counter
Applicative Counter =>
(forall a b. Counter a -> (a -> Counter b) -> Counter b)
-> (forall a b. Counter a -> Counter b -> Counter b)
-> (forall a. a -> Counter a)
-> Monad Counter
forall a. a -> Counter a
forall a b. Counter a -> Counter b -> Counter b
forall a b. Counter a -> (a -> Counter b) -> Counter b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. Counter a -> (a -> Counter b) -> Counter b
>>= :: forall a b. Counter a -> (a -> Counter b) -> Counter b
$c>> :: forall a b. Counter a -> Counter b -> Counter b
>> :: forall a b. Counter a -> Counter b -> Counter b
$creturn :: forall a. a -> Counter a
return :: forall a. a -> Counter a
Monad, Applicative Counter
Applicative Counter =>
(forall a. Counter a)
-> (forall a. Counter a -> Counter a -> Counter a)
-> (forall a. Counter a -> Counter [a])
-> (forall a. Counter a -> Counter [a])
-> Alternative Counter
forall a. Counter a
forall a. Counter a -> Counter [a]
forall a. Counter a -> Counter a -> Counter a
forall (f :: * -> *).
Applicative f =>
(forall a. f a)
-> (forall a. f a -> f a -> f a)
-> (forall a. f a -> f [a])
-> (forall a. f a -> f [a])
-> Alternative f
$cempty :: forall a. Counter a
empty :: forall a. Counter a
$c<|> :: forall a. Counter a -> Counter a -> Counter a
<|> :: forall a. Counter a -> Counter a -> Counter a
$csome :: forall a. Counter a -> Counter [a]
some :: forall a. Counter a -> Counter [a]
$cmany :: forall a. Counter a -> Counter [a]
many :: forall a. Counter a -> Counter [a]
Alternative)

-- | Creates a 'Counter' computation that represents all the keys of a 'Map'
-- occurring as many times as indicated by that key's value.
--
-- @
-- toCounter (Map.singleton k 1) === pure k
-- @
toCounter :: Map a Int -> Counter a
toCounter :: forall a. Map a Int -> Counter a
toCounter = [(a, Int)] -> Counter a
forall a b. Coercible a b => a -> b
coerce ([(a, Int)] -> Counter a)
-> (Map a Int -> [(a, Int)]) -> Map a Int -> Counter a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a Int -> [(a, Int)]
forall k a. Map k a -> [(k, a)]
Map.assocs

-- | Run a 'Counter' accumulating all the path counts into a single 'Map'
fromCounter :: Ord a => Counter a -> Map a Int
fromCounter :: forall a. Ord a => Counter a -> Map a Int
fromCounter = (Int -> Int -> Int) -> [(a, Int)] -> Map a Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([(a, Int)] -> Map a Int)
-> (Counter a -> [(a, Int)]) -> Counter a -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Counter a -> [(a, Int)]
forall a b. Coercible a b => a -> b
coerce