{-# Language BlockArguments, ImportQualifiedPost, NumericUnderscores, QuasiQuotes #-}
{-|
Module      : Main
Description : Day 23 solution
Copyright   : (c) Eric Mertens, 2020
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2020/day/23>

-}
module Main (main) where

import Advent.Format (format)
import Control.Monad (unless, zipWithM_)
import Data.Array.IO (IOUArray, getBounds, newArray_, readArray, writeArray)
import Data.Char (digitToInt)
import Data.Foldable (for_)

-- | The array maps cup numbers (indexes) to the next cup
-- in the sequence (elements).
type Ring = IOUArray Int Int

-- | Construct a /ring of cups/ given an initial arrangement and filled to
-- the given size.
newRing ::
  Int   {- ^ ring size           -} ->
  [Int] {- ^ initial arrangement -} ->
  IO Ring
newRing :: Int -> [Int] -> IO Ring
newRing Int
n [Int]
order =
  do Ring
a <- (Int, Int) -> IO Ring
forall i. Ix i => (i, i) -> IO (IOUArray i Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
newArray_ (Int
1,Int
n)
     [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
1..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] \Int
i -> Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a Int
i (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
     (Int -> Int -> IO ()) -> [Int] -> [Int] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a) [Int]
order ([Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
tail [Int]
order)
     if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
order
        then Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a ([Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
order) ([Int] -> Int
forall a. HasCallStack => [a] -> a
head [Int]
order)
        else Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a Int
n ([Int] -> Int
forall a. HasCallStack => [a] -> a
head [Int]
order) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a ([Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
order) (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+[Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int]
order)
     Ring -> IO Ring
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ring
a

readRing ::
  Ring ->
  Int {- ^ length       -} ->
  Int {- ^ starting cup -} ->
  IO [Int]
readRing :: Ring -> Int -> Int -> IO [Int]
readRing Ring
a Int
n Int
i
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Int] -> IO [Int]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
  | Bool
otherwise =
    do Int
j    <- Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
i
       [Int]
rest <- Ring -> Int -> Int -> IO [Int]
readRing Ring
a (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
j
       [Int] -> IO [Int]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
j Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
rest)

-- |
-- >>> :main
-- 47382659
-- 42271866720
main :: IO ()
IO ()
main =
  do [Int]
inp <- (Char -> Int) -> [Char] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Int
digitToInt ([Char] -> [Int]) -> IO [Char] -> IO [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [format|2020 23 %s%n|]
     [Int] -> IO ()
p1 [Int]
inp
     [Int] -> IO ()
p2 [Int]
inp

play ::
  Ring ->
  Int {- ^ iterations -} ->
  Int {- ^ current cup -} ->
  IO ()
play :: Ring -> Int -> Int -> IO ()
play Ring
a Int
i Int
cur =
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
iInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0)
  do -- extract a group of three cups
     Int
g1 <- Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
cur
     Int
g2 <- Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
g1
     Int
g3 <- Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
g2

     -- find next cup and link current one to it
     Int
nx <- Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
g3
     Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a Int
cur Int
nx

     -- find the new destination label
     (Int
_,Int
hi) <- Ring -> IO (Int, Int)
forall i. Ix i => IOUArray i Int -> IO (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds Ring
a
     let dec :: Int -> Int
dec Int
1 = Int
hi
         dec Int
x = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
         dest :: Int
dest = (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (\Int
x -> Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
g1 Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
g2 Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
g3) Int -> Int
dec (Int -> Int
dec Int
cur)

     -- splice the group back in at dest
     Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a Int
g3 (Int -> IO ()) -> IO Int -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
dest
     Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a Int
dest Int
g1

     Ring -> Int -> Int -> IO ()
play Ring
a (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
nx

p1 :: [Int] -> IO ()
p1 :: [Int] -> IO ()
p1 [Int]
inp =
  do Ring
ring <- Int -> [Int] -> IO Ring
newRing ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inp) [Int]
inp
     let sz :: Int
sz = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inp
     Ring -> Int -> Int -> IO ()
play Ring
ring Int
100 ([Int] -> Int
forall a. HasCallStack => [a] -> a
head [Int]
inp)

     [Int]
xs <- Ring -> Int -> Int -> IO [Int]
readRing Ring
ring (Int
szInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
1
     [Char] -> IO ()
putStrLn ((Int -> [Char]) -> [Int] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Int -> [Char]
forall a. Show a => a -> [Char]
show [Int]
xs)

p2 :: [Int] -> IO ()
p2 :: [Int] -> IO ()
p2 [Int]
inp =
  do let sz :: Int
sz   =  Int
1_000_000
         iter :: Int
iter = Int
10_000_000

     Ring
ring <- Int -> [Int] -> IO Ring
newRing Int
sz [Int]
inp
     Ring -> Int -> Int -> IO ()
play Ring
ring Int
iter ([Int] -> Int
forall a. HasCallStack => [a] -> a
head [Int]
inp)

     [Int]
xs <- Ring -> Int -> Int -> IO [Int]
readRing Ring
ring Int
2 Int
1
     Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
xs)