advent21/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/21/advent-of-code-2022-day-21/
   2
   3 import AoC
   4 import Data.Text (Text)
   5 import qualified Data.Text.IO as TIO
   6 import Data.Attoparsec.Text hiding (take, D)
   7 import Control.Applicative
   8 import qualified Data.Map.Strict as M
   9 import Data.Map.Strict ((!))
  10 import Control.Lens hiding (op)
  11
  12 data Shout = Literal Int | Operation Operator String String
  13   deriving (Show, Eq, Ord)
  14
  15 data Operator = Plus | Minus | Times | Divide
  16   deriving (Show, Eq, Ord)
  17
  18 type Monkeys = M.Map String Shout
  19
  20 fromLiteral :: Shout -> Int
  21 fromLiteral (Literal n) = n
  22 fromLiteral _ = error "fromLiteral"
  23
  24 main :: IO ()
  25 main =
  26   do  dataFileName <- getDataFileName
  27       text <- TIO.readFile dataFileName
  28       let monkeys = successfulParse text
  29       print $ part1 monkeys
  30       print $ part2 monkeys
  31
  32 part1, part2 :: Monkeys -> Int
  33 part1 monkeys = fromLiteral $ findRoot values operations
  34   where (values, operations) = splitMonkeys monkeys
  35
  36 -- hardcoded assumption: small values of humn give large values of result, so find larger values to give negative output
  37 part2 monkeys = binarySearch values operations l u
  38   where (Operation _ rootL rootR) = monkeys ! "root"
  39         monkeys' = monkeys & at "root" ?~ (Operation Minus rootL rootR)
  40         (values, operations) = splitMonkeys monkeys'
  41         (l, u) = findRange values operations 1
  42
  43 findRange :: Monkeys -> Monkeys -> Int -> (Int, Int)
  44 findRange values operations prev
  45   | res > 0 = findRange values operations (prev * 2)
  46   | otherwise = (prev, prev * 2)
  47   where res = fromLiteral $ trial values operations (prev * 2)
  48
  49 binarySearch :: Monkeys -> Monkeys -> Int -> Int -> Int
  50 binarySearch values operations lower upper
  51   | lower > upper = error "Failed search"
  52   | result == 0 = probe
  53   | result > 0 = binarySearch values operations (probe + 1) upper
  54   | result < 0 = binarySearch values operations lower probe
  55   where probe = ((upper - lower) `div` 2) + lower
  56         result = fromLiteral $ trial values operations probe
  57
  58 trial :: Monkeys -> Monkeys -> Int -> Shout
  59 trial values operations humn = findRoot (values & at "humn" ?~ (Literal humn)) operations
  60
  61 splitMonkeys :: Monkeys -> (Monkeys, Monkeys)
  62 splitMonkeys = M.partition f
  63   where f (Literal _) = True
  64         f (Operation _ _ _) = False
  65
  66 findRoot :: Monkeys -> Monkeys -> Shout
  67 findRoot values operations
  68   | "root" `M.member` values = values ! "root"
  69   | otherwise = findRoot values' operations'
  70   where (values', operations') = evaluateMonkeys values operations
  71
  72 evaluateMonkeys :: Monkeys -> Monkeys -> (Monkeys, Monkeys)
  73 evaluateMonkeys values operations = M.foldlWithKey' f ((values, M.empty)) operations
  74   where f (valMs, opMs) name op =
  75           case (evalShout valMs op) of
  76             Nothing -> (valMs, opMs & at name ?~ op)
  77             Just v -> (valMs & at name ?~ v, sans name opMs)
  78
  79 evalShout :: Monkeys -> Shout -> Maybe Shout
  80 evalShout _ (Literal n) = Just $ Literal n
  81 evalShout values (Operation op l r) = apply <$> (Just op) <*> lval <*> rval
  82   where lval = M.lookup l values
  83         rval = M.lookup r values
  84
  85 apply :: Operator -> Shout -> Shout -> Shout
  86 apply Plus (Literal l) (Literal r) = Literal (l + r)
  87 apply Minus (Literal l) (Literal r) = Literal (l - r)
  88 apply Times (Literal l) (Literal r) = Literal (l * r)
  89 apply Divide (Literal l) (Literal r) = Literal (l `div` r)
  90 apply _ _ _ = error "Illegal apply"
  91
  92
  93 -- Parse the input file
  94
  95 monkeysP :: Parser Monkeys
  96 monkeyP :: Parser (String, Shout)
  97 shoutP, numberP, operationP :: Parser Shout
  98 nameP :: Parser String
  99 operatorP, plusP, minusP, timesP, divideP :: Parser Operator
 100
 101 monkeysP = M.fromList <$> monkeyP `sepBy` endOfLine
 102 monkeyP = (,) <$> (nameP <* ": ") <*> shoutP
 103 shoutP = numberP <|> operationP
 104 numberP = Literal <$> decimal
 105 operationP = opify <$> nameP <*> operatorP <*> nameP
 106   where opify l o r = Operation o l r
 107 nameP = many1 letter
 108 operatorP = plusP <|> minusP <|> timesP <|> divideP
 109 plusP = Plus   <$ " + "
 110 minusP = Minus  <$ " - "
 111 timesP = Times  <$ " * "
 112 divideP = Divide <$ " / "
 113
 114 successfulParse :: Text -> Monkeys
 115 successfulParse input =
 116   case parseOnly monkeysP input of
 117     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
 118     Right monkeys -> monkeys