advent18/Main.hs

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