advent11/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/14/advent-of-code-2022-day-11/
   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 Data.List
   9 import qualified Data.IntMap.Strict as M
  10 import Data.IntMap ((!))
  11 import Control.Lens
  12 import Control.Monad.State.Strict
  13 import Control.Monad.Reader
  14 import Control.Monad.Writer
  15 import Control.Monad.RWS.Strict
  16
  17 data MonkeyCode = MonkeyCode
  18   { _operation :: Expression
  19   , _test :: Int
  20   , _trueTarget :: Int
  21   , _falseTarget :: Int
  22   }
  23   deriving (Show, Eq)
  24
  25 data Expression = Expression Operator Operand deriving (Show, Eq)
  26 data Operator = Plus | Times deriving (Show, Eq)
  27 data Operand = Literal Int | Old deriving (Show, Eq)
  28
  29 type MonkeyCodes = M.IntMap MonkeyCode
  30 data MonkeyDescription = MonkeyDescription { _limit :: Int -> Int
  31                                            , _codes :: MonkeyCodes
  32                                            }
  33 type MonkeyHolds = M.IntMap [Int]
  34 data MonkeyLog = MonkeyLog Int Int -- monkey ID, number of items handled this round
  35   deriving (Show, Eq)
  36
  37 type MonkeyHandler = RWS MonkeyDescription [MonkeyLog] MonkeyHolds
  38
  39 makeLenses ''MonkeyCode
  40
  41 main :: IO ()
  42 main =
  43   do  dataFileName <- getDataFileName
  44       text <- TIO.readFile dataFileName
  45       let (monkeyCode, monkeyHold) = successfulParse text
  46       print $ part1 monkeyCode monkeyHold
  47       print $ part2 monkeyCode monkeyHold
  48
  49 part1, part2 :: MonkeyCodes -> MonkeyHolds -> Int
  50 part1 monkeyCode monkeyHold = monkeyBusinessLevel logs
  51   where monkeyDesc = MonkeyDescription { _limit = (`div` 3)
  52                                        , _codes = monkeyCode
  53                                        }
  54         (_, logs) =  execRWS (replicateM_ 20 throwRound)
  55                              monkeyDesc monkeyHold
  56
  57 part2 monkeyCode monkeyHold = monkeyBusinessLevel logs
  58   where monkeyDesc = MonkeyDescription { _limit = (`mod` threshold)
  59                                        , _codes = monkeyCode
  60                                        }
  61         (_, logs) =  execRWS (replicateM_ 10000 throwRound)
  62                              monkeyDesc monkeyHold
  63         threshold = product $ monkeyCode ^.. folded . test
  64
  65 throwRound :: MonkeyHandler ()
  66 throwRound =
  67   do mIds <- gets M.keys
  68      mapM_ throwItems mIds
  69
  70 throwItems :: Int -> MonkeyHandler ()
  71 throwItems mId =
  72   do items <- gets (! mId)
  73      mapM_ (throwItem mId) items
  74      modify (M.insert mId [])
  75      tell [MonkeyLog mId (length items)]
  76
  77 throwItem :: Int -> Int -> MonkeyHandler ()
  78 throwItem mId currentWorry =
  79   do  monkey <- asks ((! mId) . _codes)
  80       threshold <- asks _limit
  81       let newWorry = updateWorry currentWorry (monkey ^. operation) threshold
  82       let testResult = worryTest (monkey ^. test) newWorry
  83       let recipient = if testResult
  84                       then (monkey ^. trueTarget)
  85                       else (monkey ^. falseTarget)
  86       modify (receivesItem recipient newWorry)
  87
  88 updateWorry :: Int -> Expression -> (Int -> Int) -> Int
  89 updateWorry current (Expression operator operand) threshold
  90   | operator == Plus  = threshold (current + n)
  91   | operator == Times = threshold (current * n)
  92   where n = evalOperand operand
  93         evalOperand (Literal k) = k
  94         evalOperand Old = current
  95
  96 worryTest :: Int -> Int -> Bool
  97 worryTest divisor worry = worry `mod` divisor == 0
  98
  99 receivesItem :: Int -> Int -> MonkeyHolds -> MonkeyHolds
 100 receivesItem mId worry items = M.adjust (++ [worry]) mId items
 101
 102 sumLogs :: [MonkeyLog] -> M.IntMap Int
 103 sumLogs logs = foldl' addCount M.empty logs
 104   where addCount m (MonkeyLog mId n)
 105           | mId `M.member` m = M.adjust (+ n) mId m
 106           | otherwise = M.insert mId n m
 107
 108 monkeyBusinessLevel :: [MonkeyLog] -> Int
 109 monkeyBusinessLevel logs = prolifics!!0 * prolifics!!1
 110   where prolifics = reverse $ sort $ M.elems $ sumLogs logs
 111
 112 -- Parse the input file
 113
 114 monkeysP :: Parser (MonkeyCodes, MonkeyHolds)
 115 monkeysP = makeMonkeyMaps <$> (monkeyP `sepBy` (endOfLine <* endOfLine))
 116   where makeMonkeyMaps monkeys =
 117           ( M.fromList $ map fst monkeys
 118           , M.fromList $ map snd monkeys
 119           )
 120
 121 monkeyP :: Parser ((Int, MonkeyCode), (Int, [Int]))
 122 monkeyP = mkMonkeyPair <$> mIdP <*> startingP <*> operatorP
 123                        <*> testP <*> trueTargetP <*> falseTargetP
 124   where mkMonkeyPair mId holding _operation _test _trueTarget _falseTarget =
 125           ((mId, MonkeyCode{..}), (mId, holding))
 126
 127 mIdP, testP, trueTargetP, falseTargetP :: Parser Int
 128 startingP :: Parser [Int]
 129 operatorP, expressionP :: Parser Expression
 130 opP :: Parser Operator
 131 operandP :: Parser Operand
 132
 133 mIdP = ("Monkey " *> decimal) <* ":" <* endOfLine
 134 startingP = ("  Starting items: " *> (decimal `sepBy` ", ")) <* endOfLine
 135 operatorP = ("  Operation: new = old " *> expressionP) <* endOfLine
 136 testP = ("  Test: divisible by " *> decimal) <* endOfLine
 137 trueTargetP = ("    If true: throw to monkey " *> decimal) <* endOfLine
 138 falseTargetP = ("    If false: throw to monkey " *> decimal)
 139
 140 expressionP = Expression <$> (opP <* " ") <*> operandP
 141 opP = (Plus <$ "+") <|> (Times <$ "*")
 142 operandP = (Literal <$> decimal) <|> (Old <$ "old")
 143
 144 successfulParse :: Text -> (MonkeyCodes, MonkeyHolds)
 145 successfulParse input =
 146   case parseOnly monkeysP input of
 147     Left  _err -> (M.empty, M.empty) -- TIO.putStr $ T.pack $ parseErrorPretty err
 148     Right monkeys -> monkeys