Zippers
The idea of a zipper is to represent some data structure together with a "point of interest" or "focus". The zipper makes it easy to manipulate the focus element and move around the structure to relative positions. For a tree, the zipped tree is a pair of (current element, surrounding tree) and we have the functions left, right, up (to move around the tree), and modify (to change the current element). I took the implementations of these from the Haskell wiki page.
The data structure and parser is pretty much determined by the problem, so I won't comment on them much here. Note the bespoke show instance for numbers.
data Tree = Pair Tree Tree | Leaf Int
deriving (Eq)
instance Show Tree where
show (Leaf n) = show n
show (Pair l r) = "[" ++ show l ++ "," ++ show r ++ "]"
sfNumbersP = sfNumberP `sepBy` endOfLine
sfNumberP = regularP <|> pairP
regularP = Leaf <$> decimal
pairP = Pair <$> ("[" *> sfNumberP) <*> ("," *> sfNumberP) <* "]"
(I referred to /u/ThreeFx's post while writing my solution, even if I did end up veering quite a way away from that one.)
Splitting
I'll do splitting first, as that's easier. The idea is to find the leftmost splittable element in the tree, then modify the tree to replace the element with a pair.
First is to find a splittable element. This is a function from a Tree to a Maybe Loc, returning Nothing if there are no splittable numbers. splitable wraps the tree in a context then recurses into it with splittableC.
If that's given a Pair, it looks for splittable numbers in the branches, exploiting that Maybe is Applicative so <|> handles the alternatives nicely.
splittable :: Tree -> Maybe Loc
splittable t = splittableC (top t)
splittableC :: Loc -> Maybe Loc
splittableC t@(Leaf n, _)
| n >= 10 = Just t
| otherwise = Nothing
splittableC t@(Pair _ _, _) = splittableC (left t) <|> splittableC (right t)
Given the ability to find splittable numbers, split will split the leftmost one. If there is no splittable number, it returns Nothing. If there is a splittable number, it focuses on that number, replaces it with the pair, then returns focus to the root of the tree.
split :: Tree -> Maybe Tree
split num =
case mn0 of
Nothing -> Nothing
Just _ -> Just num1
where
mn0 = splittable num
n0 = fromJust mn0
((Leaf sn), _) = n0
ln = sn `div` 2
rn = ln + sn `mod` 2
n1 = modify n0 (\_ -> Pair (Leaf ln) (Leaf rn))
(num1, _) = upmost n1
Exploding
Exploding a number is much more involved.
Finding the pair to split is much the same as above, but using a counter to track now many layers deep to find the pair.
pairAtDepth :: Int -> Tree -> Maybe Loc
pairAtDepth n t = pairAtDepthC n (top t)
pairAtDepthC :: Int -> Loc -> Maybe Loc
pairAtDepthC _ (Leaf _, _) = Nothing
pairAtDepthC 0 t@(Pair _ _, _) = Just t
pairAtDepthC n t@(Pair _ _, _) =
pairAtDepthC (n - 1) (left t) <|> pairAtDepthC (n - 1) (right t)
The complications come with finding the tree elements that are modified during the explosion.
Given a pair that explodes, I need to find the rightmost leaf that's to the left of this pair. I find that by going back up the tree until I find where I took the right branch at a pair. I then take the left branch of that pair and follow it down, always taking the right branch. If I get to the root of the tree without finding branch to the left, I return Nothing.
rightmostOnLeft :: Loc -> Maybe Loc
rightmostOnLeft (_, Top) = Nothing
rightmostOnLeft t@(_, L c r) = rightmostOnLeft $ up t
rightmostOnLeft t@(_, R l c) = Just $ rightmostNum $ left $ up t
rightmostNum :: Loc -> Loc
rightmostNum t@(Leaf _, _) = t
rightmostNum t@(Pair _ _, _) = rightmostNum $ right t
Finding the leftmost number on the right is the same, but with left and right swapped.
Actually handling the explosion is complicated, due to the fact that there may not be receiving numbers on either side. I tried using a do block to chain together all the Maybes, but that just meant that one missing number led to the whole explosion failing. The closest I got was a use of >>= to string together a couple of functions of type Loc -> Maybe Loc.
I also have to re-find the exploding pair after every update, to ensure that all the changes follow through to the end, hence the p0, p1, p2, p3. I could have reified the path itself, with functions to extract a path from one structure and reapply it to another, but this approach was good enough.
explode :: Tree -> Maybe Tree
explode num =
case mp0 of
Nothing -> Nothing
Just _ -> Just num1
where
mp0 = pairAtDepth 4 num
p0 = fromJust mp0
((Pair (Leaf nl) (Leaf nr)), _) = p0
p1 = case rightmostOnLeft p0 of
Nothing -> p0
Just leftReg -> modify leftReg (\(Leaf n) -> Leaf (n + nl))
p2 = case pairAtDepthC 4 (upmost p1) >>= leftmostOnRight of
Nothing -> p1
Just rightReg -> modify rightReg (\(Leaf n) -> Leaf (n + nr))
p3 = case pairAtDepthC 4 (upmost p2) of
Nothing -> p2
Just centrePair -> modify centrePair (\_ -> Leaf 0)
(num1, _) = upmost p3
After that, reduce is the repeated application of explode and split. Again, I use that Maybe is Applicative to simplify handling the priority of exploding over splitting.
reduce :: Tree -> Tree
reduce num = case explode num <|> split num of
Nothing -> num
Just num1 -> reduce num1
snailAdd :: Tree -> Tree -> Tree
snailAdd a b = reduce $ Pair a b
Solving the puzzle
Given all that, solving the actual puzzle is very simple!
Finding the magnitude is a simple tree walk. Adding all the numbers is a fold, finding the maximal pair is just that.
part1 numbers = magnitude total
where total = foldl1' snailAdd numbers
part2 numbers = maximum [ magnitude $ snailAdd a b
| a <- numbers, b <- numbers]
magnitude :: Tree -> Int
magnitude (Leaf n) = n
magnitude (Pair a b) = 3 * (magnitude a) + 2 * (magnitude b)
Code
You can get the code from my locally-hosed Git repo, or from Gitlab.