52 lines
1.3 KiB
Haskell
52 lines
1.3 KiB
Haskell
import Common
|
|
import qualified Data.Map as M
|
|
import Data.Function.Memoize (Memoizable(memoize))
|
|
|
|
type Monkey = String
|
|
|
|
data Operation = Const Int | Add Monkey Monkey | Sub Monkey Monkey | Mult Monkey Monkey | Div Monkey Monkey
|
|
deriving (Show, Eq)
|
|
|
|
charToOp '+' = Add
|
|
charToOp '-' = Sub
|
|
charToOp '*' = Mult
|
|
charToOp '/' = Div
|
|
charToOp c = error $ "unknown op " ++ [c]
|
|
|
|
parseOp :: String -> Operation
|
|
parseOp str
|
|
| length s == 1 = Const $ read str
|
|
| length s == 3 = op (s !! 0) (s !! 2)
|
|
| otherwise = error $ "can't parse op " ++ str
|
|
where
|
|
s = words str
|
|
op = charToOp $ head $ s !! 1
|
|
|
|
parseLine :: String -> (Monkey, Operation)
|
|
parseLine s = (name, op)
|
|
where
|
|
op = parseOp $ trim $ tail opS
|
|
(name, opS) = splitAt 4 s
|
|
|
|
rootMonkey :: Monkey
|
|
rootMonkey = "root"
|
|
|
|
eval :: M.Map Monkey Operation -> Int
|
|
eval monkeys = evalMonkey' rootMonkey
|
|
where
|
|
evalMonkey' = memoize evalMonkey
|
|
|
|
evalMonkey :: String -> Int
|
|
evalMonkey m = case op of
|
|
Const n -> n
|
|
Add a b -> evalMonkey' a + evalMonkey' b
|
|
Sub a b -> evalMonkey' a - evalMonkey' b
|
|
Mult a b -> evalMonkey' a * evalMonkey' b
|
|
Div a b -> evalMonkey' a `div` evalMonkey' b
|
|
where
|
|
op = (M.!) monkeys m
|
|
|
|
main = interact $
|
|
show
|
|
. eval
|
|
. M.fromList . map parseLine . lines |