Done day 21

[advent-of-code-22.git] / advent21 / Main.hs

diff --git a/advent21/Main.hs b/advent21/Main.hs
new file mode 100644 (file)
index 0000000..1811612
--- /dev/null
+++ b/advent21/Main.hs
@@ -0,0 +1,118 @@
+-- Writeup at https://work.njae.me.uk/2022/12/21/advent-of-code-2022-day-21/
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+import Control.Lens
+
+data Shout = Literal Int | Operation Operator String String
+  deriving (Show, Eq, Ord)
+
+data Operator = Plus | Minus | Times | Divide 
+  deriving (Show, Eq, Ord)
+
+type Monkeys = M.Map String Shout
+
+fromLiteral :: Shout -> Int
+fromLiteral (Literal n) = n
+fromLiteral _ = error "fromLiteral"
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let monkeys = successfulParse text
+      print $ part1 monkeys
+      print $ part2 monkeys
+
+part1, part2 :: Monkeys -> Int
+part1 monkeys = fromLiteral $ findRoot values operations
+  where (values, operations) = splitMonkeys monkeys
+
+-- hardcoded assumption: small values of humn give large values of result, so find larger values to give negative output
+part2 monkeys = binarySearch values operations l u
+  where (Operation _ rootL rootR) = monkeys ! "root"
+        monkeys' = monkeys & at "root" ?~ (Operation Minus rootL rootR)
+        (values, operations) = splitMonkeys monkeys'
+        (l, u) = findRange values operations 1
+
+findRange :: Monkeys -> Monkeys -> Int -> (Int, Int)
+findRange values operations prev 
+  | res > 0 = findRange values operations (prev * 2)
+  | otherwise = (prev, prev * 2)
+  where res = fromLiteral $ trial values operations (prev * 2)
+
+binarySearch :: Monkeys -> Monkeys -> Int -> Int -> Int
+binarySearch values operations lower upper 
+  | lower > upper = error "Failed search"
+  | result == 0 = probe
+  | result > 0 = binarySearch values operations (probe + 1) upper
+  | result < 0 = binarySearch values operations lower probe
+  where probe = ((upper - lower) `div` 2) + lower
+        result = fromLiteral $ trial values operations probe
+
+trial :: Monkeys -> Monkeys -> Int -> Shout
+trial values operations humn = findRoot (values & at "humn" ?~ (Literal humn)) operations
+
+splitMonkeys :: Monkeys -> (Monkeys, Monkeys)
+splitMonkeys = M.partition f
+  where f (Literal _) = True
+        f (Operation _ _ _) = False
+
+findRoot :: Monkeys -> Monkeys -> Shout
+findRoot values operations 
+  | "root" `M.member` values = values ! "root"
+  | otherwise = findRoot values' operations'
+  where (values', operations') = evaluateMonkeys values operations
+
+evaluateMonkeys :: Monkeys -> Monkeys -> (Monkeys, Monkeys)
+evaluateMonkeys values operations = M.foldlWithKey' f ((values, M.empty)) operations
+  where f (valMs, opMs) name op = 
+          case (evalShout valMs op) of
+            Nothing -> (valMs, opMs & at name ?~ op)
+            Just v -> (valMs & at name ?~ v, sans name opMs)
+
+evalShout :: Monkeys -> Shout -> Maybe Shout
+evalShout _ (Literal n) = Just $ Literal n
+evalShout values (Operation op l r) = apply <$> (Just op) <*> lval <*> rval
+  where lval = M.lookup l values
+        rval = M.lookup r values
+
+apply :: Operator -> Shout -> Shout -> Shout
+apply Plus (Literal l) (Literal r) = Literal (l + r)
+apply Minus (Literal l) (Literal r) = Literal (l - r)
+apply Times (Literal l) (Literal r) = Literal (l * r)
+apply Divide (Literal l) (Literal r) = Literal (l `div` r)
+apply _ _ _ = error "Illegal apply"
+
+
+-- Parse the input file
+
+monkeysP :: Parser Monkeys
+monkeyP :: Parser (String, Shout)
+shoutP, numberP, operationP :: Parser Shout
+nameP :: Parser String
+operatorP, plusP, minusP, timesP, divideP :: Parser Operator
+
+monkeysP = M.fromList <$> monkeyP `sepBy` endOfLine
+monkeyP = (,) <$> (nameP <* ": ") <*> shoutP
+shoutP = numberP <|> operationP
+numberP = Literal <$> decimal
+operationP = opify <$> nameP <*> operatorP <*> nameP
+  where opify l o r = Operation o l r
+nameP = many1 letter
+operatorP = plusP <|> minusP <|> timesP <|> divideP
+plusP = Plus   <$ " + "
+minusP = Minus  <$ " - "
+timesP = Times  <$ " * "
+divideP = Divide <$ " / "
+
+successfulParse :: Text -> Monkeys
+successfulParse input = 
+  case parseOnly monkeysP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right monkeys -> monkeys