Tidying
[advent-of-code-22.git] / 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 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