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

<https://adventofcode.com/2015/day/20>

Run a simulation of elves delivering presents which each elf taking a larger
step size than the previous.

-}
module Main where

import Advent.Format (format)
import Control.Monad.Loop (for, exec_)
import Control.Monad.Trans.Class ( MonadTrans(lift) )
import Data.Array.ST (readArray, writeArray, MArray(newArray), runSTUArray)
import Data.Array.Unboxed (UArray, assocs)

-- | >>> :main
-- 831600
-- 884520
main :: IO ()
IO ()
main =
 do Int
target <- [format|2015 20 %u%n|]
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> UArray Int Int -> Int
findHouse Int
target (Int -> UArray Int Int
solve1 Int
target))
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> UArray Int Int -> Int
findHouse Int
target (Int -> UArray Int Int
solve2 Int
target))

-- | Return the house number with at least the given number of presents.
findHouse :: Int -> UArray Int Int -> Int
findHouse :: Int -> UArray Int Int -> Int
findHouse Int
target UArray Int Int
a = [Int] -> Int
forall a. HasCallStack => [a] -> a
head [Int
h | (Int
h,Int
t) <- UArray Int Int -> [(Int, Int)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Int Int
a, Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
target]

solve1 :: Int -> UArray Int Int
solve1 :: Int -> UArray Int Int
solve1 Int
target = (forall s. ST s (STUArray s Int Int)) -> UArray Int Int
forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray
 do let top :: Int
top = Int
target Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
10
    STUArray s Int Int
a <- (Int, Int) -> Int -> ST s (STUArray s Int Int)
forall i. Ix i => (i, i) -> Int -> ST s (STUArray s i Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1,Int
top) Int
0
    STUArray s Int Int
a STUArray s Int Int -> ST s () -> ST s (STUArray s Int Int)
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ LoopT (ST s) () -> ST s ()
forall (m :: * -> *) a. Applicative m => LoopT m a -> m ()
exec_
     do Int
elf   <- Int -> (Int -> Bool) -> (Int -> Int) -> LoopT (ST s) Int
forall a (m :: * -> *). a -> (a -> Bool) -> (a -> a) -> LoopT m a
for   Int
1 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
top) (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+)
        Int
house <- Int -> (Int -> Bool) -> (Int -> Int) -> LoopT (ST s) Int
forall a (m :: * -> *). a -> (a -> Bool) -> (a -> a) -> LoopT m a
for Int
elf (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
top) (Int
elfInt -> Int -> Int
forall a. Num a => a -> a -> a
+)
        ST s () -> LoopT (ST s) ()
forall (m :: * -> *) a. Monad m => m a -> LoopT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift do Int
old <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
a Int
house
                STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
a Int
house (Int
old Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
elfInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
10)

solve2 :: Int -> UArray Int Int
solve2 :: Int -> UArray Int Int
solve2 Int
target = (forall s. ST s (STUArray s Int Int)) -> UArray Int Int
forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray
 do let top :: Int
top = Int
target Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
11
    STUArray s Int Int
a <- (Int, Int) -> Int -> ST s (STUArray s Int Int)
forall i. Ix i => (i, i) -> Int -> ST s (STUArray s i Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1,Int
top) Int
0
    STUArray s Int Int
a STUArray s Int Int -> ST s () -> ST s (STUArray s Int Int)
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ LoopT (ST s) () -> ST s ()
forall (m :: * -> *) a. Applicative m => LoopT m a -> m ()
exec_
     do Int
elf   <- Int -> (Int -> Bool) -> (Int -> Int) -> LoopT (ST s) Int
forall a (m :: * -> *). a -> (a -> Bool) -> (a -> a) -> LoopT m a
for   Int
1 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<=Int
top) (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+)
        Int
house <- Int -> (Int -> Bool) -> (Int -> Int) -> LoopT (ST s) Int
forall a (m :: * -> *). a -> (a -> Bool) -> (a -> a) -> LoopT m a
for Int
elf (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
top (Int
elfInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
50)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
elf)
        ST s () -> LoopT (ST s) ()
forall (m :: * -> *) a. Monad m => m a -> LoopT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift do Int
old <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
a Int
house
                STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
a Int
house (Int
old Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
elfInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
11)