93 lines
2.7 KiB
Haskell
93 lines
2.7 KiB
Haskell
import Common
|
|
import qualified Data.Map as M
|
|
import Data.Function.Memoize (Memoizable(memoize))
|
|
import Data.Maybe (fromJust)
|
|
|
|
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]
|
|
|
|
getBranches :: Operation -> Maybe (Monkey, Monkey)
|
|
getBranches (Const _) = Nothing
|
|
getBranches (Add a b) = Just (a, b)
|
|
getBranches (Sub a b) = Just (a, b)
|
|
getBranches (Mult a b) = Just (a, b)
|
|
getBranches (Div a b) = Just (a, b)
|
|
|
|
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"
|
|
human :: Monkey
|
|
human = "humn"
|
|
|
|
eval :: M.Map Monkey Operation -> Int
|
|
eval monkeys = getInput humanBranch targetResult
|
|
where
|
|
getInput m n = case monkey of
|
|
-- a + b = n
|
|
Add a b -> if lHuman then getInput a (n - evalMonkey' b) else getInput b (n - evalMonkey' a)
|
|
-- a - b = n
|
|
Sub a b -> if lHuman then getInput a (n + evalMonkey' b) else getInput b (evalMonkey' a - n)
|
|
-- a * b = n
|
|
Mult a b -> if lHuman then getInput a (n `div` evalMonkey' b) else getInput b (n `div` evalMonkey' a)
|
|
-- a / b = n
|
|
Div a b -> if lHuman then getInput a (n * evalMonkey' b) else getInput b (evalMonkey' a `div` n)
|
|
Const _ -> n
|
|
where
|
|
lHuman = hasHuman' branchL
|
|
(branchL, branchR) = fromJust $ getBranches $ (M.!) monkeys m
|
|
monkey = (M.!) monkeys m
|
|
|
|
targetResult = evalMonkey' resolvedBranch
|
|
|
|
resolvedBranch = if hasHuman' rootBranchL then rootBranchR else rootBranchL
|
|
humanBranch = if hasHuman' rootBranchL then rootBranchL else rootBranchR
|
|
|
|
(rootBranchL, rootBranchR) = fromJust $ getBranches $ (M.!) monkeys rootMonkey
|
|
|
|
hasHuman' = memoize hasHuman
|
|
|
|
hasHuman :: Monkey -> Bool
|
|
hasHuman m
|
|
| m == human = True
|
|
| otherwise = case getBranches $ (M.!) monkeys m of
|
|
Just (a, b) -> hasHuman' a || hasHuman' b
|
|
Nothing -> False
|
|
|
|
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 |