91 lines
3.2 KiB
Haskell
91 lines
3.2 KiB
Haskell
import Data.List.Split (splitOn)
|
|
import qualified Control.Monad.Trans.State as St
|
|
import qualified Data.Map as M
|
|
import Control.Monad (replicateM_)
|
|
import Data.List (sort)
|
|
|
|
data Expr = Old | Lit Int | Add Expr Expr | Mul Expr Expr
|
|
data Test = T Int Int Int
|
|
data Part = Part1 | Part2
|
|
data Monkey = M [Int] Expr Test
|
|
|
|
getFactor :: Monkey -> Int
|
|
getFactor (M _ _ (T x _ _)) = x
|
|
|
|
parseElem :: String -> Expr
|
|
parseElem "old" = Old
|
|
parseElem i = Lit (read i)
|
|
|
|
parseExpr :: [String] -> Expr
|
|
parseExpr s
|
|
| s !! 1 == "*" = Mul (parseElem $ s !! 0) (parseElem $ s !! 2)
|
|
| s !! 1 == "+" = Add (parseElem $ s !! 0) (parseElem $ s !! 2)
|
|
| otherwise = error $ "unknown expression: " ++ concat s
|
|
|
|
parseMonkey :: [String] -> Monkey
|
|
parseMonkey s = M startingItems expr test where
|
|
startingItems = read $ "[" ++ (splitOn ": " (s !! 1) !! 1) ++ "]"
|
|
expr = parseExpr $ last $ splitOn ["="] (words (s !! 2))
|
|
testFactor = read $ last $ words (s !! 3)
|
|
testTargetT = read $ last $ words (s !! 4)
|
|
testTargetF = read $ last $ words (s !! 5)
|
|
test = T testFactor testTargetT testTargetF
|
|
|
|
evalExpr :: Expr -> Int -> Int
|
|
evalExpr Old x = x
|
|
evalExpr (Lit i) x = i
|
|
evalExpr (Add a b) x = evalExpr a x + evalExpr b x
|
|
evalExpr (Mul a b) x = evalExpr a x * evalExpr b x
|
|
|
|
type Monkeys = M.Map Int (Monkey, Int)
|
|
|
|
steps :: Part -> Int -> Monkey -> St.State Monkeys (Int, Monkey)
|
|
steps part factor monkey@(M [] _ _) = do
|
|
pure (0, monkey)
|
|
steps part factor (M (item:items) op test@(T fac tr fa)) = do
|
|
let newItem' = evalExpr op item --`div` 3
|
|
let newItem = case part of
|
|
Part1 -> evalExpr op item `div` 3
|
|
Part2 -> evalExpr op item `mod` factor
|
|
let targetMonkey = if newItem `mod` fac == 0 then tr else fa
|
|
monkeys <- St.get
|
|
let (M items' op' test', count') = monkeys M.! targetMonkey
|
|
let updated = M (newItem:items') op' test'
|
|
St.put $ M.insert targetMonkey (updated, count') monkeys
|
|
(count, final) <- steps part factor $ M items op test
|
|
pure (count+1, final)
|
|
|
|
simulate :: Part -> Int -> Int -> St.State Monkeys ()
|
|
simulate part factor i = do
|
|
monkeys <- St.get
|
|
let (monkey, count) = monkeys M.! i
|
|
(iters, finished) <- steps part factor monkey
|
|
monkeys <- St.get
|
|
St.put $ M.insert i (finished, count + iters) monkeys
|
|
|
|
round :: Part -> Int -> St.State Monkeys ()
|
|
round part factor = do
|
|
monkeys <- St.get
|
|
mapM_ (simulate part factor . fst) (M.toAscList monkeys)
|
|
|
|
main :: IO ()
|
|
main = do
|
|
monkeys <- map parseMonkey . splitOn [""] . lines <$> getContents
|
|
let factor = foldr1 lcm $ map getFactor monkeys
|
|
let initialState = M.fromAscList $ zip [0..] (map (\x -> (x,0)) monkeys)
|
|
|
|
putStr "part 1: "
|
|
let action = replicateM_ 20 (Main.round Part1 factor)
|
|
let results = St.execState action initialState
|
|
let counts = map (snd . snd) (M.toAscList results)
|
|
print $ counts
|
|
print $ product $ take 2 $ reverse $ sort counts
|
|
|
|
putStr "part 2: "
|
|
let action = replicateM_ 10000 (Main.round Part2 factor)
|
|
let results = St.execState action initialState
|
|
let counts = map (snd . snd) (M.toAscList results)
|
|
print $ counts
|
|
print $ product $ take 2 $ reverse $ sort counts
|
|
|