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

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

Today's problem had us perform manipulations on a tree-based
term language. It was made tricky because the problem asked
us to do things to the nearest left and right neighbors
of elements of our tree.

-}
module Main (main) where

import Advent (format, stageTH)
import Control.Applicative ((<|>))
import Data.List (foldl1', tails)
import Text.ParserCombinators.ReadP

-- * Binary trees

-- | A binary tree with integers at the leaves
data Tree
  = Tree :+: Tree -- ^ tuple
  | Leaf !Int  -- ^ number
  deriving Int -> Tree -> ShowS
[Tree] -> ShowS
Tree -> String
(Int -> Tree -> ShowS)
-> (Tree -> String) -> ([Tree] -> ShowS) -> Show Tree
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Tree -> ShowS
showsPrec :: Int -> Tree -> ShowS
$cshow :: Tree -> String
show :: Tree -> String
$cshowList :: [Tree] -> ShowS
showList :: [Tree] -> ShowS
Show

-- * Parsing

-- | Tree parser from a leaf parser
t :: ReadP Tree
t :: ReadP Tree
t =
  Int -> Tree
Leaf (Int -> Tree) -> ReadP Int -> ReadP Tree
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReadS Int -> ReadP Int
forall a. ReadS a -> ReadP a
readS_to_P ReadS Int
forall a. Read a => ReadS a
reads ReadP Tree -> ReadP Tree -> ReadP Tree
forall a. ReadP a -> ReadP a -> ReadP a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|>
  Tree -> Tree -> Tree
(:+:) (Tree -> Tree -> Tree)
-> ReadP String -> ReadP (Tree -> Tree -> Tree)
forall a b. a -> ReadP b -> ReadP a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ ReadP String
"[" ReadP (Tree -> Tree -> Tree) -> ReadP Tree -> ReadP (Tree -> Tree)
forall a b. ReadP (a -> b) -> ReadP a -> ReadP b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReadP Tree
t ReadP (Tree -> Tree) -> ReadP String -> ReadP (Tree -> Tree)
forall a b. ReadP a -> ReadP b -> ReadP a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ReadP String
"," ReadP (Tree -> Tree) -> ReadP Tree -> ReadP Tree
forall a b. ReadP (a -> b) -> ReadP a -> ReadP b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReadP Tree
t ReadP Tree -> ReadP String -> ReadP Tree
forall a b. ReadP a -> ReadP b -> ReadP a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ReadP String
"]"

stageTH

-- | >>> :main
-- 3551
-- 4555
main :: IO ()
IO ()
main =
 do [Tree]
inp <- [format|2021 18 (@t%n)*|]
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Tree -> Int
magnitude ((Tree -> Tree -> Tree) -> [Tree] -> Tree
forall a. HasCallStack => (a -> a -> a) -> [a] -> a
foldl1' Tree -> Tree -> Tree
add [Tree]
inp))
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Tree -> Int
magnitude (Tree -> Tree -> Tree
add Tree
x Tree
y) Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Tree -> Int
magnitude (Tree -> Tree -> Tree
add Tree
y Tree
x)
                   | Tree
x:[Tree]
ys <- [Tree] -> [[Tree]]
forall a. [a] -> [[a]]
tails [Tree]
inp, Tree
y <- [Tree]
ys])

-- * Snailfish operations

-- | Add two expressions and reduce them
add :: Tree -> Tree -> Tree
add :: Tree -> Tree -> Tree
add Tree
x Tree
y = Tree -> Tree
reduce (Tree
x Tree -> Tree -> Tree
:+: Tree
y)

-- | Reduce an expression until it won't reduce
reduce :: Tree -> Tree
reduce :: Tree -> Tree
reduce (Tree -> Tree
explode -> Tree
x) = Tree -> (Tree -> Tree) -> Maybe Tree -> Tree
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Tree
x Tree -> Tree
reduce (Tree -> Maybe Tree
split Tree
x) 

-- | Explode /all/ the pairs at depth 4 from left to right.
explode :: Tree -> Tree
explode :: Tree -> Tree
explode = Int -> [(Side, Tree)] -> Tree -> Tree
forall {t}. (Num t, Ord t) => t -> [(Side, Tree)] -> Tree -> Tree
down (Int
0::Int) []
  where
    down :: t -> [(Side, Tree)] -> Tree -> Tree
down t
4 [(Side, Tree)]
z (Leaf Int
l :+: Leaf Int
r) = t -> [(Side, Tree)] -> Tree -> Tree
up t
4 (Side -> Int -> [(Side, Tree)] -> [(Side, Tree)]
sendUp Side
L Int
l (Side -> Int -> [(Side, Tree)] -> [(Side, Tree)]
sendUp Side
R Int
r [(Side, Tree)]
z)) (Int -> Tree
Leaf Int
0)
    down t
d [(Side, Tree)]
z (Tree
l :+: Tree
r) | t
d t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
4   = t -> [(Side, Tree)] -> Tree -> Tree
down (t
dt -> t -> t
forall a. Num a => a -> a -> a
+t
1) ((Side
R,Tree
r)(Side, Tree) -> [(Side, Tree)] -> [(Side, Tree)]
forall a. a -> [a] -> [a]
:[(Side, Tree)]
z) Tree
l
    down t
d [(Side, Tree)]
z Tree
x                   = t -> [(Side, Tree)] -> Tree -> Tree
up t
d [(Side, Tree)]
z Tree
x

    up :: t -> [(Side, Tree)] -> Tree -> Tree
up t
_ []        Tree
x = Tree
x
    up t
d ((Side
R,Tree
r):[(Side, Tree)]
z) Tree
l = t -> [(Side, Tree)] -> Tree -> Tree
down t
d ((Side
L,Tree
l)(Side, Tree) -> [(Side, Tree)] -> [(Side, Tree)]
forall a. a -> [a] -> [a]
:[(Side, Tree)]
z) Tree
r
    up t
d ((Side
L,Tree
l):[(Side, Tree)]
z) Tree
r = t -> [(Side, Tree)] -> Tree -> Tree
up (t
dt -> t -> t
forall a. Num a => a -> a -> a
-t
1) [(Side, Tree)]
z (Tree
l Tree -> Tree -> Tree
:+: Tree
r)

-- | Replace the first number with value 10 or more with a pair
-- of it divided in half rounding first down then up.
split :: Tree -> Maybe Tree
split :: Tree -> Maybe Tree
split (Tree
l :+: Tree
r) = (Tree -> Tree -> Tree
:+: Tree
r) (Tree -> Tree) -> Maybe Tree -> Maybe Tree
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree -> Maybe Tree
split Tree
l Maybe Tree -> Maybe Tree -> Maybe Tree
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (Tree
l Tree -> Tree -> Tree
:+:) (Tree -> Tree) -> Maybe Tree -> Maybe Tree
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree -> Maybe Tree
split Tree
r
split (Leaf Int
x)
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
10, (Int
q,Int
r) <- Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
quotRem Int
x Int
2 = Tree -> Maybe Tree
forall a. a -> Maybe a
Just (Int -> Tree
Leaf Int
q Tree -> Tree -> Tree
:+: Int -> Tree
Leaf (Int
qInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r))
  | Bool
otherwise                     = Maybe Tree
forall a. Maybe a
Nothing

-- | Compute the /magnitude/ of an expression
--
-- >>> parse = fst . head . readP_to_S t
--
-- >>> magnitude (parse "[9,1]")
-- 29
--
-- >>> magnitude (parse "[[1,2],[[3,4],5]]")
-- 143
magnitude :: Tree -> Int
magnitude :: Tree -> Int
magnitude (Leaf Int
x) = Int
x
magnitude (Tree
l :+: Tree
r) = Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Tree -> Int
magnitude Tree
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Tree -> Int
magnitude Tree
r

-- * Tree zippers

-- | Marks the side of a tree node constructor that
-- we know the subtree of.
data Side = L | R deriving (Side -> Side -> Bool
(Side -> Side -> Bool) -> (Side -> Side -> Bool) -> Eq Side
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Side -> Side -> Bool
== :: Side -> Side -> Bool
$c/= :: Side -> Side -> Bool
/= :: Side -> Side -> Bool
Eq, Int -> Side -> ShowS
[Side] -> ShowS
Side -> String
(Int -> Side -> ShowS)
-> (Side -> String) -> ([Side] -> ShowS) -> Show Side
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Side -> ShowS
showsPrec :: Int -> Side -> ShowS
$cshow :: Side -> String
show :: Side -> String
$cshowList :: [Side] -> ShowS
showList :: [Side] -> ShowS
Show)

-- | A hole in a binary tree.
type Zip = [(Side, Tree)]

-- | Add a number to the nearest sibling on the given side.
sendUp :: Side -> Int -> Zip -> Zip
sendUp :: Side -> Int -> [(Side, Tree)] -> [(Side, Tree)]
sendUp Side
L Int
x ((Side
L,Tree
l):[(Side, Tree)]
zs) = (Side
L, Int -> Tree -> Tree
sendR Int
x Tree
l)(Side, Tree) -> [(Side, Tree)] -> [(Side, Tree)]
forall a. a -> [a] -> [a]
:[(Side, Tree)]
zs
sendUp Side
R Int
x ((Side
R,Tree
r):[(Side, Tree)]
zs) = (Side
R, Int -> Tree -> Tree
sendL Int
x Tree
r)(Side, Tree) -> [(Side, Tree)] -> [(Side, Tree)]
forall a. a -> [a] -> [a]
:[(Side, Tree)]
zs
sendUp Side
h Int
x ((Side, Tree)
z    :[(Side, Tree)]
zs) = (Side, Tree)
z (Side, Tree) -> [(Side, Tree)] -> [(Side, Tree)]
forall a. a -> [a] -> [a]
: Side -> Int -> [(Side, Tree)] -> [(Side, Tree)]
sendUp Side
h Int
x [(Side, Tree)]
zs
sendUp Side
_ Int
_ []         = []

-- | Add a number to the left-most leaf
sendL :: Int -> Tree -> Tree
sendL :: Int -> Tree -> Tree
sendL Int
x (Tree
l :+: Tree
r) = Int -> Tree -> Tree
sendL Int
x Tree
l Tree -> Tree -> Tree
:+: Tree
r
sendL Int
x (Leaf Int
y)  = Int -> Tree
Leaf (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y)

-- | Add a number to the rightmost leaf
sendR :: Int -> Tree -> Tree
sendR :: Int -> Tree -> Tree
sendR Int
x (Tree
l :+: Tree
r) = Tree
l Tree -> Tree -> Tree
:+: Int -> Tree -> Tree
sendR Int
x Tree
r
sendR Int
x (Leaf Int
y)  = Int -> Tree
Leaf (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y)