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

<https://adventofcode.com/2018/day/8>

Parse a tree out of a list of integers and then answer a pair of queries
about the tree.

-}
{-# Language DeriveTraversable #-}
module Main (main) where

import Advent (format)
import Control.Monad (replicateM)
import Control.Monad.Trans.State.Strict (StateT(..))
import Data.List (uncons)

-- | Print the answers to day 8
--
-- >>> :main
-- 42196
-- 33649
main :: IO ()
IO ()
main =
 do [Int]
input <- [format|2018 8 %u& %n|]
    let Just (Tree Int
tree, []) = StateT [Int] Maybe (Tree Int) -> [Int] -> Maybe (Tree Int, [Int])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT [Int] Maybe (Tree Int)
parseTree [Int]
input
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Tree Int -> Int
forall a. Num a => Tree a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Tree Int
tree)
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Tree Int -> Int
part2 Tree Int
tree)

-- | A tree can have children and metadata entries.
data Tree a = Tree [Tree a] [a] -- ^ children and metadata
  deriving ((forall a b. (a -> b) -> Tree a -> Tree b)
-> (forall a b. a -> Tree b -> Tree a) -> Functor Tree
forall a b. a -> Tree b -> Tree a
forall a b. (a -> b) -> Tree a -> Tree 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) -> Tree a -> Tree b
fmap :: forall a b. (a -> b) -> Tree a -> Tree b
$c<$ :: forall a b. a -> Tree b -> Tree a
<$ :: forall a b. a -> Tree b -> Tree a
Functor, (forall m. Monoid m => Tree m -> m)
-> (forall m a. Monoid m => (a -> m) -> Tree a -> m)
-> (forall m a. Monoid m => (a -> m) -> Tree a -> m)
-> (forall a b. (a -> b -> b) -> b -> Tree a -> b)
-> (forall a b. (a -> b -> b) -> b -> Tree a -> b)
-> (forall b a. (b -> a -> b) -> b -> Tree a -> b)
-> (forall b a. (b -> a -> b) -> b -> Tree a -> b)
-> (forall a. (a -> a -> a) -> Tree a -> a)
-> (forall a. (a -> a -> a) -> Tree a -> a)
-> (forall a. Tree a -> [a])
-> (forall a. Tree a -> Bool)
-> (forall a. Tree a -> Int)
-> (forall a. Eq a => a -> Tree a -> Bool)
-> (forall a. Ord a => Tree a -> a)
-> (forall a. Ord a => Tree a -> a)
-> (forall a. Num a => Tree a -> a)
-> (forall a. Num a => Tree a -> a)
-> Foldable Tree
forall a. Eq a => a -> Tree a -> Bool
forall a. Num a => Tree a -> a
forall a. Ord a => Tree a -> a
forall m. Monoid m => Tree m -> m
forall a. Tree a -> Bool
forall a. Tree a -> Int
forall a. Tree a -> [a]
forall a. (a -> a -> a) -> Tree a -> a
forall m a. Monoid m => (a -> m) -> Tree a -> m
forall b a. (b -> a -> b) -> b -> Tree a -> b
forall a b. (a -> b -> b) -> b -> Tree a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => Tree m -> m
fold :: forall m. Monoid m => Tree m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Tree a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Tree a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Tree a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Tree a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> Tree a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Tree a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Tree a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Tree a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Tree a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Tree a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Tree a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Tree a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> Tree a -> a
foldr1 :: forall a. (a -> a -> a) -> Tree a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Tree a -> a
foldl1 :: forall a. (a -> a -> a) -> Tree a -> a
$ctoList :: forall a. Tree a -> [a]
toList :: forall a. Tree a -> [a]
$cnull :: forall a. Tree a -> Bool
null :: forall a. Tree a -> Bool
$clength :: forall a. Tree a -> Int
length :: forall a. Tree a -> Int
$celem :: forall a. Eq a => a -> Tree a -> Bool
elem :: forall a. Eq a => a -> Tree a -> Bool
$cmaximum :: forall a. Ord a => Tree a -> a
maximum :: forall a. Ord a => Tree a -> a
$cminimum :: forall a. Ord a => Tree a -> a
minimum :: forall a. Ord a => Tree a -> a
$csum :: forall a. Num a => Tree a -> a
sum :: forall a. Num a => Tree a -> a
$cproduct :: forall a. Num a => Tree a -> a
product :: forall a. Num a => Tree a -> a
Foldable, Functor Tree
Foldable Tree
(Functor Tree, Foldable Tree) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Tree a -> f (Tree b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Tree (f a) -> f (Tree a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Tree a -> m (Tree b))
-> (forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a))
-> Traversable Tree
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
sequenceA :: forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
$csequence :: forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
sequence :: forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
Traversable, Int -> Tree a -> ShowS
[Tree a] -> ShowS
Tree a -> String
(Int -> Tree a -> ShowS)
-> (Tree a -> String) -> ([Tree a] -> ShowS) -> Show (Tree a)
forall a. Show a => Int -> Tree a -> ShowS
forall a. Show a => [Tree a] -> ShowS
forall a. Show a => Tree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
showsPrec :: Int -> Tree a -> ShowS
$cshow :: forall a. Show a => Tree a -> String
show :: Tree a -> String
$cshowList :: forall a. Show a => [Tree a] -> ShowS
showList :: [Tree a] -> ShowS
Show)

-- | Sum of metadata entries on leaf nodes and recursive call on
-- child nodes identified by indexes stored in metadata.
part2 :: Tree Int -> Int
part2 :: Tree Int -> Int
part2 (Tree [Tree Int]
xs [Int]
ys)
  | [Tree Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Tree Int]
xs   = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
ys
  | Bool
otherwise = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Maybe Int -> Int
forall a. Num a => Maybe a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int -> Maybe Int
forall a. [a] -> Int -> Maybe a
index ((Tree Int -> Int) -> [Tree Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Tree Int -> Int
part2 [Tree Int]
xs) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) | Int
i <- [Int]
ys, Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0]

-- | list index returning Nothing on failure.
index :: [a] -> Int -> Maybe a
index :: forall a. [a] -> Int -> Maybe a
index [a]
xs Int
n
  | a
a:[a]
_ <- Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
n [a]
xs = a -> Maybe a
forall a. a -> Maybe a
Just a
a
  | Bool
otherwise        = Maybe a
forall a. Maybe a
Nothing

-- | Get the next integer
int :: StateT [Int] Maybe Int
int :: StateT [Int] Maybe Int
int = ([Int] -> Maybe (Int, [Int])) -> StateT [Int] Maybe Int
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT [Int] -> Maybe (Int, [Int])
forall a. [a] -> Maybe (a, [a])
uncons

-- | Parse a tree from a list of integers
parseTree :: StateT [Int] Maybe (Tree Int)
parseTree :: StateT [Int] Maybe (Tree Int)
parseTree =
 do Int
n <- StateT [Int] Maybe Int
int
    Int
m <- StateT [Int] Maybe Int
int
    [Tree Int]
a <- Int
-> StateT [Int] Maybe (Tree Int) -> StateT [Int] Maybe [Tree Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n StateT [Int] Maybe (Tree Int)
parseTree
    [Int]
b <- Int -> StateT [Int] Maybe Int -> StateT [Int] Maybe [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
m StateT [Int] Maybe Int
int
    Tree Int -> StateT [Int] Maybe (Tree Int)
forall a. a -> StateT [Int] Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Tree Int] -> [Int] -> Tree Int
forall a. [Tree a] -> [a] -> Tree a
Tree [Tree Int]
a [Int]
b)