Learning Python with Advent of Code Walkthroughs

Dazbo's Advent of Code solutions, written in Python

Monkeys

Advent of Code 2022 - Day 21

Day 21: Monkey Math

Useful Links

Concepts and Packages Demonstrated

regexrecursionoperatorassertionBinary search

Problem Intro

Part 1 was pretty easy. Hurrah! Part 2… Required a bit more thinking!

We’re told that if we solve the monkeys’ riddle, they will show us a shortcut to the grove. The riddle is in the form of a set of instructions that look like this:

root: pppw + sjmn
dbpl: 5
cczh: sllz + lgvd
zczc: 2
ptdq: humn - dvpt
dvpt: 3
lfqf: 4
humn: 5
ljgn: 2
sjmn: drzm * dbpl
sllz: 4
pppw: cczh / lfqf
lgvd: ljgn * ptdq
drzm: hmdt - zczc
hmdt: 32

The item before the : is the monkey name. Each monkey has one job. The job is one of:

The math operations are always one of +, -, * or /.

Part 1

What number will the monkey named root yell?

This is a recursive problem. We either know the numeric value the monkey needs to yell, or we need to recursively perform a math operation between two other monkeys.

Let’s read in the data…

    with open(INPUT_FILE, mode="rt") as f:
        data = f.read().splitlines()

    # first, let's get all the monkeys with known yell values
    yell_pattern = re.compile(r"([a-z]{4}): (\d+)")
    calc_pattern = re.compile(r"([a-z]{4}): ([a-z]{4}) (.){1,2} ([a-z]{4})")

    # Part 1    
    monkeys: dict[str, int] = {}  # monkey: value    
    rows_to_remove = []
    for row, line in enumerate(data):
        if match := yell_pattern.match(line):
            monkeys[match.groups()[0]] = int(match.groups()[1])
            rows_to_remove.append(row)

    # strip out the known monkeys, and leave only the calculations
    calcs_only = [line for row, line in enumerate(data) if row not in rows_to_remove]
    
    calcs = {} # Assemble a dict of {monkey_id: (monkey, op, monkey), ...}
    for line in calcs_only:
        monkey_id, monkey2, op, monkey3 = calc_pattern.findall(line)[0]
        calcs[monkey_id] = (monkey2, op, monkey3)

First I’m going through every line and identifying only monkeys where we know the numeric yell value. For each one:

With the sample data, our monkeys now looks like this:

{'dbpl': 5, 'zczc': 2, 'dvpt': 3, 'lfqf': 4, 'humn': 5, 'ljgn': 2, 'sllz': 4, 'hmdt': 32}

All the rows that remain in the input data must be operations with a pair of monkeys. I use list comprehension to obtain a new list of all the input rows that are not enumerated in our rows_to_remove list. This gives us a list that I’ve called calcs_only. For each of these, I create an entry in a calcs dict, where the value is a tuple containing a pair of monkeys, and the operation to perform between them.

With the same data, calcs now looks like this:

root: ('pppw', '+', 'sjmn')
cczh: ('sllz', '+', 'lgvd')
ptdq: ('humn', '-', 'dvpt')
sjmn: ('drzm', '*', 'dbpl')
pppw: ('cczh', '/', 'lfqf')
lgvd: ('ljgn', '*', 'ptdq')
drzm: ('hmdt', '-', 'zczc')

Now I create a dictionary that maps a given operator to an operator method from the operator module. (Don’t forget to import operator.)

OPMAP = {
    '+': operator.add,
    '*': operator.mul,
    '-': operator.sub,
    '/': operator.floordiv
}

Now I create a recursive function that evaluates the value of any monkey specified by ID:

def evaluate_monkey(monkey_id: str, calcs, monkeys) -> int:
    """ Recursive evaluation of calcs like: pppw + sjmn """
    assert monkey_id in calcs, "We require a monkey with an operation"
    current_calc = calcs[monkey_id]
    monkey2, monkey3 = current_calc[0], current_calc[2]
    op = current_calc[1]
    
    # recurse for monkeys we don't yet know value for
    if monkey2 not in monkeys:
        evaluate_monkey(monkey2, calcs, monkeys)
    if monkey3 not in monkeys:
        evaluate_monkey(monkey3, calcs, monkeys)

    # base case
    # We could use eval, but it's dangerous, and relatively slow
    # monkeys[monkey_id] = int(eval(str(monkeys[monkey2]) + op + str(monkeys[monkey3])))
    monkeys[monkey_id] = OPMAP[op](monkeys[monkey2], monkeys[monkey3])
    
    return monkeys[monkey_id]

Here’s how it works:

Finally, we can use our recursive function to evaluate the root monkey:

    evaluate_monkey("root", calcs, monkeys)
    print(f"Part 1: root={monkeys['root']}")

Easy so far!

Part 2

It turns out that the root monkey doesn’t add up two numbers. Instead, this monkey is performing an equality check: i.e. whether the two values in the operation match.

Furthermore, it turns out that the monkey named humn isn’t a monkey at all. It’s me! The existing value for this entry is irrelevant. My job is to find the number that must be yelled in order for root’s equality test to be True.

What number do you yell to pass root’s equality test?

My strategy:

Here’s the code to recreate our known monkeys, and to update the operation for the root monkey:

    monkeys = {}  # recreate known monkeys, {monkey: value}
    for row, line in enumerate(data):
        if match := yell_pattern.match(line):
            monkeys[match.groups()[0]] = int(match.groups()[1])

    # change the root monkey instruction. We'll change it to a subtract operator.
    # That way, we'll know both operands have the same value when the result is 0
    calcs["root"] = (calcs["root"][0], "-", calcs["root"][2])

And now for how we “guess” the yell value of humn

I started by just performing a linear search, staring at 0 and incrementing by one with each try. But this was far too slow to solve the problem. However, it was useful because I was able to see that as my value of humn was increased, so too was the final value for root. So, given this relationship, it stands to reason that a binary search should be effective.

Here’s how I’ve implemented the search:

def binary_search(target, low:int, high:int, func, *func_args, reverse_search=False) -> int:
    """ Generic binary search function that takes a target to find,
    low and high values to start with, and a function to run, plus its args. 
    Implicitly returns None if the search is exceeded. """
    
    res = None  # just set it to something that isn't the target
    candidate = 0  # initialise; we'll set it to the mid point in a second
    
    while low < high:  # search exceeded        
        candidate = int((low+high) // 2)  # pick mid-point of our low and high        
        # print(f"{candidate}->{res}")
        res = func(candidate, *func_args) # run our function, whatever it is
        if res == target:
            return candidate  # solution found
        
        if res > target:
            low = candidate
        else:
            high = candidate

def try_monkeys(candidate, calcs: dict, monkeys: dict) -> int:
    monkeys_try = monkeys.copy()
    monkeys_try["humn"] = candidate
    res = evaluate_monkey("root", calcs, monkeys_try)
    return res

The try_monkeys() function applies a candidate value for humn, applies the instructions, and determines the resulting value of root. Recall that the goal is for this to be 0.

In our binary_search(), we define low and high bounds to try for humn, and then invoke the try_monkeys() function with the midpoint of these values. If the result is not our goal, then we split our search range in half, and find the midpoint of the appropriate half. Here, if the result is larger than our goal, then we set the low value to be our previous midpoint; thus, we’re now only searching the top half of the original range.

Interestingly, I found that my binary search worked for the sample data, but not for the real data. That’s because the relationship between the size of humn and root was reversed in my real data! So I modified my binary search so that we can optionally reverse the direction of the > test. If the first search fails to find a valid number, then we simply re-run the search, but with the comparison operator reversed.

I.e.

    # Part 2
    monkeys = {}  # recreate known monkeys, {monkey: value}
    for row, line in enumerate(data):
        if match := yell_pattern.match(line):
            monkeys[match.groups()[0]] = int(match.groups()[1])

    # change the root monkey instruction. We'll change it to a subtract operator.
    # That way, we'll know both operands have the same value when the result is 0
    calcs["root"] = (calcs["root"][0], "-", calcs["root"][2])
    
    # We need to try values that will return will result in root == 0.
    # Brute force is really slow, but binary search works well!
    humn = binary_search(0, 0, 1e16, try_monkeys, calcs, monkeys)
    if humn is None: # try reverse correlation binary search
        humn = binary_search(0, 0, 1e16, try_monkeys, calcs, monkeys, reverse_search=True)
    print(f"Part 2: humn={humn}")

def binary_search(target, low:int, high:int, func, *func_args, reverse_search=False) -> int:
    """ Generic binary search function that takes a target to find,
    low and high values to start with, and a function to run, plus its args. 
    Implicitly returns None if the search is exceeded. """
    
    res = None  # just set it to something that isn't the target
    candidate = 0  # initialise; we'll set it to the mid point in a second
    
    while low < high:  # search exceeded        
        candidate = int((low+high) // 2)  # pick mid-point of our low and high        
        # print(f"{candidate}->{res}")
        res = func(candidate, *func_args) # run our function, whatever it is
        if res == target:
            return candidate  # solution found
        
        comp = operator.gt if not reverse_search else operator.lt
        if comp(res, target):
            low = candidate
        else:
            high = candidate

Results

Here’s the final code:

import operator
from pathlib import Path
import re
import time

SCRIPT_DIR = Path(__file__).parent
# INPUT_FILE = Path(SCRIPT_DIR, "input/sample_input.txt")
INPUT_FILE = Path(SCRIPT_DIR, "input/input.txt")

OPMAP = {
    '+': operator.add,
    '*': operator.mul,
    '-': operator.sub,
    '/': operator.floordiv
}

def main():
    with open(INPUT_FILE, mode="rt") as f:
        data = f.read().splitlines()

    # first, let's get all the monkeys with known yell values
    yell_pattern = re.compile(r"([a-z]{4}): (\d+)")
    calc_pattern = re.compile(r"([a-z]{4}): ([a-z]{4}) (.){1,2} ([a-z]{4})")

    # Part 1    
    monkeys: dict[str, int] = {}  # monkey: value    
    rows_to_remove = []
    for row, line in enumerate(data):
        if match := yell_pattern.match(line):
            monkeys[match.groups()[0]] = int(match.groups()[1])
            rows_to_remove.append(row)

    # strip out the known monkeys, and leave only the calculations
    calcs_only = [line for row, line in enumerate(data) if row not in rows_to_remove]
    
    calcs = {} # Assemble a dict of {monkey_id: (monkey, op, monkey), ...}
    for line in calcs_only:
        monkey_id, monkey2, op, monkey3 = calc_pattern.findall(line)[0]
        calcs[monkey_id] = (monkey2, op, monkey3)

    evaluate_monkey("root", calcs, monkeys)
    print(f"Part 1: root={monkeys['root']}")
    
    # Part 2
    monkeys = {}  # recreate known monkeys, {monkey: value}
    for row, line in enumerate(data):
        if match := yell_pattern.match(line):
            monkeys[match.groups()[0]] = int(match.groups()[1])

    # change the root monkey instruction. We'll change it to a subtract operator.
    # That way, we'll know both operands have the same value when the result is 0
    calcs["root"] = (calcs["root"][0], "-", calcs["root"][2])
    
    # We need to try values that will return will result in root == 0.
    # Brute force is really slow, but binary search works well!
    humn = binary_search(0, 0, 1e16, try_monkeys, calcs, monkeys)
    if humn is None: # try reverse correlation binary search
        humn = binary_search(0, 0, 1e16, try_monkeys, calcs, monkeys, reverse_search=True)
    print(f"Part 2: humn={humn}")
    
def binary_search(target, low:int, high:int, func, *func_args, reverse_search=False) -> int:
    """ Generic binary search function that takes a target to find,
    low and high values to start with, and a function to run, plus its args. 
    Implicitly returns None if the search is exceeded. """
    
    res = None  # just set it to something that isn't the target
    candidate = 0  # initialise; we'll set it to the mid point in a second
    
    while low < high:  # search exceeded        
        candidate = int((low+high) // 2)  # pick mid-point of our low and high        
        # print(f"{candidate}->{res}")
        res = func(candidate, *func_args) # run our function, whatever it is
        if res == target:
            return candidate  # solution found
        
        comp = operator.gt if not reverse_search else operator.lt
        if comp(res, target):
            low = candidate
        else:
            high = candidate

def try_monkeys(candidate, calcs: dict, monkeys: dict) -> int:
    monkeys_try = monkeys.copy()
    monkeys_try["humn"] = candidate
    res = evaluate_monkey("root", calcs, monkeys_try)
    return res
    
def evaluate_monkey(monkey_id: str, calcs, monkeys) -> int:
    """ Recursive evaluation of calcs like: pppw + sjmn """
    current_calc = calcs[monkey_id]
    monkey2, monkey3 = current_calc[0], current_calc[2]
    op = current_calc[1]
    
    # recurse for monkeys we don't yet know value for
    if monkey2 not in monkeys:
        evaluate_monkey(monkey2, calcs, monkeys)
    if monkey3 not in monkeys:
        evaluate_monkey(monkey3, calcs, monkeys)

    # base case
    # We could use eval, but it's dangerous, and relatively slow
    # monkeys[monkey_id] = int(eval(str(monkeys[monkey2]) + op + str(monkeys[monkey3])))
    monkeys[monkey_id] = OPMAP[op](monkeys[monkey2], monkeys[monkey3])
    
    return monkeys[monkey_id]
            
if __name__ == "__main__":
    t1 = time.perf_counter()
    main()
    t2 = time.perf_counter()
    print(f"Execution time: {t2 - t1:0.4f} seconds")

And here’s my output:

Part 1: root=168502451381566
Part 2: humn=3343167719438
Execution time: 0.0252 seconds