advent18/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/21/advent-of-code-2021-day-18/
   2
   3 import Data.Text ()
   4 import qualified Data.Text.IO as TIO
   5 import Data.Attoparsec.Text
   6 import Control.Applicative
   7 import Data.Maybe
   8 import Data.List
   9
  10 data Tree = Pair Tree Tree | Leaf Int
  11   deriving (Eq)
  12
  13 instance Show Tree where
  14   show (Leaf n) = show n
  15   show (Pair l r) = "[" ++ show l ++ "," ++ show r ++ "]"
  16
  17 data Cxt = Top | L Cxt Tree | R Tree Cxt
  18   deriving (Show, Eq)
  19
  20 -- type Context = [(Direction, Tree)]
  21 -- data Direction = Lft | Rght
  22 --   deriving (Show, Eq)
  23
  24 type Loc = (Tree, Cxt)
  25
  26 main :: IO ()
  27 main =
  28   do  text <- TIO.readFile "data/advent18.txt"
  29       let numbers = successfulParse text
  30       print $ part1 numbers
  31       print $ part2 numbers
  32
  33 part1 numbers = magnitude total
  34   where total = foldl1' snailAdd numbers
  35
  36 part2 numbers = maximum [ magnitude $ snailAdd a b
  37                         | a <- numbers, b <- numbers]
  38
  39 magnitude :: Tree -> Int
  40 magnitude (Leaf n) = n
  41 magnitude (Pair a b) = 3 * (magnitude a) + 2 * (magnitude b)
  42
  43
  44 left :: Loc -> Loc
  45 left (Pair l r, c) = (l, L c r)
  46
  47 right :: Loc -> Loc
  48 right (Pair l r, c) = (r, R l c)
  49
  50 top :: Tree -> Loc
  51 top t = (t, Top)
  52
  53 up :: Loc -> Loc
  54 up (t, Top) = (t, Top)
  55 up (t, L c r) = (Pair t r, c)
  56 up (t, R l c) = (Pair l t, c)
  57
  58 upmost :: Loc -> Loc
  59 upmost l@(t, Top) = l
  60 upmost l = upmost (up l)
  61
  62 modify :: Loc -> (Tree -> Tree) -> Loc
  63 modify (t, c) f = (f t, c)
  64
  65
  66 explode :: Tree -> Maybe Tree
  67 explode num =
  68   case mp0 of
  69     Nothing -> Nothing
  70     Just _ -> Just num1
  71   where
  72     mp0 = pairAtDepth 4 num
  73     p0 = fromJust mp0
  74     ((Pair (Leaf nl) (Leaf nr)), _) = p0
  75     p1 = case rightmostOnLeft p0 of
  76       Nothing -> p0
  77       Just leftReg -> modify leftReg (\(Leaf n) -> Leaf (n + nl))
  78     p2 = case pairAtDepthC 4 (upmost p1) >>= leftmostOnRight of
  79       Nothing -> p1
  80       Just rightReg -> modify rightReg (\(Leaf n) -> Leaf (n + nr))
  81     p3 = case pairAtDepthC 4 (upmost p2) of
  82       Nothing -> p2
  83       Just centrePair -> modify centrePair (\_ -> Leaf 0)
  84     (num1, _) = upmost p3
  85
  86 pairAtDepth :: Int -> Tree -> Maybe Loc
  87 pairAtDepth n t = pairAtDepthC n (top t)
  88
  89 pairAtDepthC :: Int -> Loc -> Maybe Loc
  90 pairAtDepthC _ (Leaf _, _) = Nothing
  91 pairAtDepthC 0 t@(Pair _ _, _) = Just t
  92 pairAtDepthC n t@(Pair _ _, _) =
  93   pairAtDepthC (n - 1) (left t) <|> pairAtDepthC (n - 1) (right t)
  94
  95 rightmostOnLeft :: Loc -> Maybe Loc
  96 rightmostOnLeft (_, Top) = Nothing
  97 rightmostOnLeft t@(_, L c r) = rightmostOnLeft $ up t
  98 rightmostOnLeft t@(_, R l c) = Just $ rightmostNum $ left $ up t
  99
 100 rightmostNum :: Loc -> Loc
 101 rightmostNum t@(Leaf _, _) = t
 102 rightmostNum t@(Pair _ _, _) = rightmostNum $ right t
 103
 104 leftmostOnRight :: Loc -> Maybe Loc
 105 leftmostOnRight (_, Top) = Nothing
 106 leftmostOnRight t@(_, R l c) = leftmostOnRight $ up t
 107 leftmostOnRight t@(_, L c r) = Just $ leftmostNum $ right $ up t
 108
 109 leftmostNum :: Loc -> Loc
 110 leftmostNum t@(Leaf _, _) = t
 111 leftmostNum t@(Pair _ _, _) = leftmostNum $ left t
 112
 113 split :: Tree -> Maybe Tree
 114 split num =
 115   case mn0 of
 116     Nothing -> Nothing
 117     Just _ -> Just num1
 118   where
 119     mn0 = splittable num
 120     n0 = fromJust mn0
 121     ((Leaf sn), _) = n0
 122     ln = sn `div` 2
 123     rn = ln + sn `mod` 2
 124     n1 = modify n0 (\_ -> Pair (Leaf ln) (Leaf rn))
 125     (num1, _) = upmost n1
 126
 127 splittable :: Tree -> Maybe Loc
 128 splittable t = splittableC (top t)
 129
 130 splittableC :: Loc -> Maybe Loc
 131 splittableC t@(Leaf n, _)
 132   | n >= 10 = Just t
 133   | otherwise = Nothing
 134 splittableC t@(Pair _ _, _) = splittableC (left t) <|> splittableC (right t)
 135
 136 reduce :: Tree -> Tree
 137 reduce num = case explode num <|> split num of
 138   Nothing -> num
 139   Just num1 -> reduce num1
 140
 141 snailAdd :: Tree -> Tree -> Tree
 142 snailAdd a b = reduce $ Pair a b
 143
 144
 145 -- Parse the input file
 146
 147 sfNumbersP = sfNumberP `sepBy` endOfLine
 148
 149 sfNumberP = regularP <|> pairP
 150
 151 regularP = Leaf <$> decimal
 152 pairP = Pair <$> ("[" *> sfNumberP) <*> ("," *> sfNumberP) <* "]"
 153
 154 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 155 successfulParse input =
 156   case parseOnly sfNumbersP input of
 157     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 158     Right numbers -> numbers
 159
 160
 161