Advent of Code 2022 day 21

Day 21 was something we've not fully seen before this year: the evaluation of expressions. It also gave me an excuse to experiment with the index lens operations like at. I'm not sure they added much to this example.

Data structures

The data was a pretty direct translation from the problem: a monkey can Shout a number, or shout some combination of other monkeys' shouts. The group of monkeys is represented as a Map from monkey name to Shout.

Parsing the monkeys is long, to account for the nested structures and options.

data Shout = Literal Int | Operation Operator String String
  deriving (Show, Eq, Ord)

data Operator = Plus | Minus | Times | Divide 
  deriving (Show, Eq, Ord)

type Monkeys = M.Map String Shout

fromLiteral :: Shout -> Int
fromLiteral (Literal n) = n
fromLiteral _ = error "fromLiteral"

monkeysP = M.fromList <$> monkeyP `sepBy` endOfLine
monkeyP = (,) <$> (nameP <* ": ") <*> shoutP
shoutP = numberP <|> operationP
numberP = Literal <$> decimal
operationP = opify <$> nameP <*> operatorP <*> nameP
  where opify l o r = Operation o l r
nameP = many1 letter
operatorP = plusP <|> minusP <|> timesP <|> divideP
plusP = Plus   <$ " + "
minusP = Minus  <$ " - "
timesP = Times  <$ " * "
divideP = Divide <$ " / "

Evaluating monkeys

The basic idea is to split the monkeys into two groups: those that can shout a literal number, and those that need to work out what to say (called values and operations below).

splitMonkeys :: Monkeys -> (Monkeys, Monkeys)
splitMonkeys = M.partition f
  where f (Literal _) = True
        f (Operation _ _ _) = False

Then, I can go through all the monkeys with operations and, one by one, check if all the monkeys they need are ones with values. If this monkey can now shout a number, I add that monkey/number to the the values and remove it from the operations. If not, that monkey is added to a new set of pending operations.

evaluateMonkeys :: Monkeys -> Monkeys -> (Monkeys, Monkeys)
evaluateMonkeys values operations = M.foldlWithKey' f ((values, M.empty)) operations
  where f (valMs, opMs) name op = 
          case (evalShout valMs op) of
            Nothing -> (valMs, opMs & at name ?~ op)
            Just v -> (valMs & at name ?~ v, sans name opMs)

I use Applicatives to find the value of a Shout, returning Nothing if the monkey can't shout yet, or Just (Literal n) if it can shout a number.

evalShout :: Monkeys -> Shout -> Maybe Shout
evalShout _ (Literal n) = Just $ Literal n
evalShout values (Operation op l r) = apply <$> (Just op) <*> lval <*> rval
  where lval = M.lookup l values
        rval = M.lookup r values

apply :: Operator -> Shout -> Shout -> Shout
apply Plus (Literal l) (Literal r) = Literal (l + r)
apply Minus (Literal l) (Literal r) = Literal (l - r)
apply Times (Literal l) (Literal r) = Literal (l * r)
apply Divide (Literal l) (Literal r) = Literal (l `div` r)
apply _ _ _ = error "Illegal apply"

I find the value of "root" by repeating evaluateMonkeys until "root" has a value. (There may be a few monkeys that will never shout a value: I don't know, I didn't check.)

findRoot :: Monkeys -> Monkeys -> Shout
findRoot values operations 
  | "root" `M.member` values = values ! "root"
  | otherwise = findRoot values' operations'
  where (values', operations') = evaluateMonkeys values operations
  
part1 :: Monkeys -> Int
part1 monkeys = fromLiteral $ findRoot values operations
  where (values, operations) = splitMonkeys monkeys

Finding what to shout

For part 2, I changed the operation of "root" to be Minus. If the two operands were equal, "root" would shout zero. I then had to find a value to make that true, using non-zero values as clues.

A bit of poking around, using the trial function, showed that small values of "humn" gave large values of "root", and increasing the values of "humn" gave smaller values of "root". I hard-coded that behaviour into the search (meaning that my code doesn't work on the example).

trial :: Monkeys -> Monkeys -> Int -> Shout
trial values operations humn = findRoot (values & at "humn" ?~ (Literal humn)) operations

First was to find a range of values of "humn" to search. I started at 1 and kept doubling it until "root" changed from positive to negative.

findRange :: Monkeys -> Monkeys -> Int -> (Int, Int)
findRange values operations prev 
  | res > 0 = findRange values operations (prev * 2)
  | otherwise = (prev, prev * 2)
  where res = fromLiteral $ trial values operations (prev * 2)

Given a range of values where the value of "root" changed from positive to negative, I then used a binary search to find the value of "humn" that gave zero for "root".

part2 monkeys = binarySearch values operations l u
  where (Operation _ rootL rootR) = monkeys ! "root"
        monkeys' = monkeys & at "root" ?~ (Operation Minus rootL rootR)
        (values, operations) = splitMonkeys monkeys'
        (l, u) = findRange values operations 1

binarySearch :: Monkeys -> Monkeys -> Int -> Int -> Int
binarySearch values operations lower upper 
  | lower > upper = error "Failed search"
  | result == 0 = probe
  | result > 0 = binarySearch values operations (probe + 1) upper
  | result < 0 = binarySearch values operations lower probe
  where probe = ((upper - lower) `div` 2) + lower
        result = fromLiteral $ trial values operations probe

Code

You can get the code from my locally-hosted Git repo, or from Gitlab.