{-# Language BlockArguments, ViewPatterns, LambdaCase #-}
{-|
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 (getInputLines)
import Control.Applicative ((<|>))
import Data.Char (isDigit)
import Data.List (tails)
import Text.ParserCombinators.ReadP (ReadP, char, munch1, readP_to_S)

-- | >>> :main
-- 3551
-- 4555
main :: IO ()
main :: IO ()
main =
 do [Tree Int]
inp <- (String -> Tree Int) -> [String] -> [Tree Int]
forall a b. (a -> b) -> [a] -> [b]
map String -> Tree Int
parse ([String] -> [Tree Int]) -> IO [String] -> IO [Tree Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO [String]
getInputLines Int
18
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Tree Int -> Int
magnitude ((Tree Int -> Tree Int -> Tree Int) -> [Tree Int] -> Tree Int
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Tree Int -> Tree Int -> Tree Int
add [Tree Int]
inp))
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Tree Int -> Int
magnitude (Tree Int -> Tree Int -> Tree Int
add Tree Int
x Tree Int
y) Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Tree Int -> Int
magnitude (Tree Int -> Tree Int -> Tree Int
add Tree Int
y Tree Int
x)
                   | Tree Int
x:[Tree Int]
ys <- [Tree Int] -> [[Tree Int]]
forall a. [a] -> [[a]]
tails [Tree Int]
inp, Tree Int
y <- [Tree Int]
ys])

-- * Snailfish operations

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

-- | Reduce an expression until it won't reduce
reduce :: Tree Int -> Tree Int
reduce :: Tree Int -> Tree Int
reduce Tree Int
x = Tree Int -> (Tree Int -> Tree Int) -> Maybe (Tree Int) -> Tree Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Tree Int
x Tree Int -> Tree Int
reduce ((Int, Int, Zip Int) -> Tree Int
explode ((Int, Int, Zip Int) -> Tree Int)
-> Maybe (Int, Int, Zip Int) -> Maybe (Tree Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree Int -> Maybe (Int, Int, Zip Int)
forall a. Tree a -> Maybe (a, a, Zip a)
unstable Tree Int
x Maybe (Tree Int) -> Maybe (Tree Int) -> Maybe (Tree Int)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Tree Int -> Maybe (Tree Int)
split Tree Int
x)

-- | Find the left-most pair at depth 4.
unstable :: Tree a -> Maybe (a, a, Zip a)
unstable :: forall a. Tree a -> Maybe (a, a, Zip a)
unstable = Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
forall a. Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
go Int
4 []
  where
    go :: Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
    go :: forall a. Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
go Int
0 Zip a
z (Leaf a
l :+: Leaf a
r) = (a, a, Zip a) -> Maybe (a, a, Zip a)
forall a. a -> Maybe a
Just (a
l, a
r, Zip a
z)
    go Int
0 Zip a
_ Tree a
_ = Maybe (a, a, Zip a)
forall a. Maybe a
Nothing
    go Int
d Zip a
z (Tree a
l :+: Tree a
r) = Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
forall a. Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
go (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ((Side
R,Tree a
r)(Side, Tree a) -> Zip a -> Zip a
forall a. a -> [a] -> [a]
:Zip a
z) Tree a
l Maybe (a, a, Zip a) -> Maybe (a, a, Zip a) -> Maybe (a, a, Zip a)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
forall a. Int -> Zip a -> Tree a -> Maybe (a, a, Zip a)
go (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ((Side
L,Tree a
l)(Side, Tree a) -> Zip a -> Zip a
forall a. a -> [a] -> [a]
:Zip a
z) Tree a
r
    go Int
_ Zip a
_ Tree a
_ = Maybe (a, a, Zip a)
forall a. Maybe a
Nothing

-- Add the left and right components to the nearest number
-- on the left and right respectively. Replace the hole with
-- a zero.
explode :: (Int, Int, Zip Int) -> Tree Int
explode :: (Int, Int, Zip Int) -> Tree Int
explode (Int
l, Int
r, Zip Int
z)
  = Tree Int -> Zip Int -> Tree Int
forall a. Tree a -> Zip a -> Tree a
fromZip (Int -> Tree Int
forall a. a -> Tree a
Leaf Int
0)
  (Zip Int -> Tree Int) -> Zip Int -> Tree Int
forall a b. (a -> b) -> a -> b
$ Side -> (Tree Int -> Tree Int) -> Zip Int -> Zip Int
forall a. Side -> (Tree a -> Tree a) -> Zip a -> Zip a
appUp Side
L ((Int -> Int) -> Tree Int -> Tree Int
forall a. (a -> a) -> Tree a -> Tree a
appR (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+))
  (Zip Int -> Zip Int) -> Zip Int -> Zip Int
forall a b. (a -> b) -> a -> b
$ Side -> (Tree Int -> Tree Int) -> Zip Int -> Zip Int
forall a. Side -> (Tree a -> Tree a) -> Zip a -> Zip a
appUp Side
R ((Int -> Int) -> Tree Int -> Tree Int
forall a. (a -> a) -> Tree a -> Tree a
appL (Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
+)) Zip Int
z

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

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

-- * Binary trees

-- | A binary tree with data at the leaves
data Tree a
  = Tree a :+: Tree a -- ^ tuple
  | Leaf a  -- ^ regular number
  deriving 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
showList :: [Tree a] -> ShowS
$cshowList :: forall a. Show a => [Tree a] -> ShowS
show :: Tree a -> String
$cshow :: forall a. Show a => Tree a -> String
showsPrec :: Int -> Tree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
Show

-- * 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
/= :: Side -> Side -> Bool
$c/= :: Side -> Side -> Bool
== :: Side -> Side -> Bool
$c== :: 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
showList :: [Side] -> ShowS
$cshowList :: [Side] -> ShowS
show :: Side -> String
$cshow :: Side -> String
showsPrec :: Int -> Side -> ShowS
$cshowsPrec :: Int -> Side -> ShowS
Show)

-- | A hole in a binary tree. Rebuild the tree with 'fromZip'
type Zip a = [(Side, Tree a)]

-- | Rebuild a tree given a zipper and the value to put in the hole.
fromZip :: Tree a -> Zip a -> Tree a
fromZip :: forall a. Tree a -> Zip a -> Tree a
fromZip = (Tree a -> (Side, Tree a) -> Tree a)
-> Tree a -> [(Side, Tree a)] -> Tree a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl \Tree a
x -> \case
  (Side
L, Tree a
l) -> Tree a
l Tree a -> Tree a -> Tree a
forall a. Tree a -> Tree a -> Tree a
:+: Tree a
x
  (Side
R, Tree a
r) -> Tree a
x Tree a -> Tree a -> Tree a
forall a. Tree a -> Tree a -> Tree a
:+: Tree a
r

-- | Apply the given function to the nearest parent
-- sibling on the given side.
appUp :: Side -> (Tree a -> Tree a) -> Zip a -> Zip a
appUp :: forall a. Side -> (Tree a -> Tree a) -> Zip a -> Zip a
appUp Side
L Tree a -> Tree a
f ((Side
L,Tree a
l):[(Side, Tree a)]
zs) = (Side
L, Tree a -> Tree a
f Tree a
l)(Side, Tree a) -> [(Side, Tree a)] -> [(Side, Tree a)]
forall a. a -> [a] -> [a]
:[(Side, Tree a)]
zs
appUp Side
R Tree a -> Tree a
f ((Side
R,Tree a
r):[(Side, Tree a)]
zs) = (Side
R, Tree a -> Tree a
f Tree a
r)(Side, Tree a) -> [(Side, Tree a)] -> [(Side, Tree a)]
forall a. a -> [a] -> [a]
:[(Side, Tree a)]
zs
appUp Side
h Tree a -> Tree a
f ((Side, Tree a)
z    :[(Side, Tree a)]
zs) = (Side, Tree a)
z (Side, Tree a) -> [(Side, Tree a)] -> [(Side, Tree a)]
forall a. a -> [a] -> [a]
: Side -> (Tree a -> Tree a) -> [(Side, Tree a)] -> [(Side, Tree a)]
forall a. Side -> (Tree a -> Tree a) -> Zip a -> Zip a
appUp Side
h Tree a -> Tree a
f [(Side, Tree a)]
zs
appUp Side
_ Tree a -> Tree a
_ []         = []

-- | Apply a function to the left-most leaf
appL :: (a -> a) -> Tree a -> Tree a
appL :: forall a. (a -> a) -> Tree a -> Tree a
appL a -> a
f (Tree a
l :+: Tree a
r) = (a -> a) -> Tree a -> Tree a
forall a. (a -> a) -> Tree a -> Tree a
appL a -> a
f Tree a
l Tree a -> Tree a -> Tree a
forall a. Tree a -> Tree a -> Tree a
:+: Tree a
r
appL a -> a
f (Leaf a
x) = a -> Tree a
forall a. a -> Tree a
Leaf (a -> a
f a
x)

-- | Apply a function to the rightmost leaf
appR :: (a -> a) -> Tree a -> Tree a
appR :: forall a. (a -> a) -> Tree a -> Tree a
appR a -> a
f (Tree a
l :+: Tree a
r) = Tree a
l Tree a -> Tree a -> Tree a
forall a. Tree a -> Tree a -> Tree a
:+: (a -> a) -> Tree a -> Tree a
forall a. (a -> a) -> Tree a -> Tree a
appR a -> a
f Tree a
r
appR a -> a
f (Leaf a
x) = a -> Tree a
forall a. a -> Tree a
Leaf (a -> a
f a
x)

-- * Parsing

-- | Parse an expression
--
-- >>> parse "[[[[0,9],2],3],4]"
-- (((Leaf 0 :+: Leaf 9) :+: Leaf 2) :+: Leaf 3) :+: Leaf 4
parse :: String -> Tree Int
parse :: String -> Tree Int
parse (ReadP (Tree Int) -> ReadS (Tree Int)
forall a. ReadP a -> ReadS a
readP_to_S (ReadP Int -> ReadP (Tree Int)
forall a. ReadP a -> ReadP (Tree a)
pTree ReadP Int
pInt) -> [(Tree Int
x,String
_)]) = Tree Int
x
parse String
_ = String -> Tree Int
forall a. HasCallStack => String -> a
error String
"bad input"

-- | Unsigned 'Int' parser
pInt :: ReadP Int
pInt :: ReadP Int
pInt = String -> Int
forall a. Read a => String -> a
read (String -> Int) -> ReadP String -> ReadP Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Char -> Bool) -> ReadP String
munch1 Char -> Bool
isDigit

-- | ReadP expression parser
pTree :: ReadP a -> ReadP (Tree a)
pTree :: forall a. ReadP a -> ReadP (Tree a)
pTree ReadP a
pLeaf = ReadP (Tree a)
tuple ReadP (Tree a) -> ReadP (Tree a) -> ReadP (Tree a)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ReadP (Tree a)
number
  where
    tuple :: ReadP (Tree a)
tuple = Tree a -> Tree a -> Tree a
forall a. Tree a -> Tree a -> Tree a
(:+:) (Tree a -> Tree a -> Tree a)
-> ReadP Char -> ReadP (Tree a -> Tree a -> Tree a)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Char -> ReadP Char
char Char
'[' ReadP (Tree a -> Tree a -> Tree a)
-> ReadP (Tree a) -> ReadP (Tree a -> Tree a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReadP a -> ReadP (Tree a)
forall a. ReadP a -> ReadP (Tree a)
pTree ReadP a
pLeaf ReadP (Tree a -> Tree a) -> ReadP Char -> ReadP (Tree a -> Tree a)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Char -> ReadP Char
char Char
',' ReadP (Tree a -> Tree a) -> ReadP (Tree a) -> ReadP (Tree a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReadP a -> ReadP (Tree a)
forall a. ReadP a -> ReadP (Tree a)
pTree ReadP a
pLeaf ReadP (Tree a) -> ReadP Char -> ReadP (Tree a)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Char -> ReadP Char
char Char
']'
    number :: ReadP (Tree a)
number = a -> Tree a
forall a. a -> Tree a
Leaf (a -> Tree a) -> ReadP a -> ReadP (Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReadP a
pLeaf