aoc-2022/21.py

from queue import deque

lines = [line.strip().split(": ") for line in open("21.input")]

def part1():
    known = {p[0]: int(p[1]) for p in lines if p[1].isdecimal()}
    unknown = deque([(p[0], p[1]) for p in lines if not p[1].isdecimal()])

    while "root" not in known:
        name, exp = unknown.popleft()
        left, op, right = exp.split(" ")
        if left in known and right in known:
            known[name] = eval(f"known[\"{left}\"] {op} known[\"{right}\"]")
        else:
            unknown.append((name, exp))

    print("Part 1:", int(known["root"]))


monkeys = {p[0]: p[1] for p in lines}

def simplify(name):
    exp = monkeys[name]

    if name == "humn":
        return "x"

    if exp.isdecimal():
        return int(exp)

    left, op, right = exp.split(" ")

    left, right = simplify(left), simplify(right)

    if isinstance(left, int) and isinstance(right, int):
        return int(eval(f"left {op} right"))
    if left == "x" or right == "x":
        return (left, op, right)

    return (left, op, right)

def solver(lhs, rhs):
    while lhs != "x":
        left, op, right = lhs
        if op == "+":
            if isinstance(left, int):
                rhs -= left
                lhs = right
                continue
            elif isinstance(right, int):
                rhs -= right
                lhs = left
                continue
        if op == "-":
            if isinstance(left, int):
                rhs = -rhs
                rhs += left
                lhs = right
                continue
            elif isinstance(right, int):
                rhs += right
                lhs = left
                continue
        if op == "/":
            if isinstance(left, int):
                raise ValueError("Cannot solve with x in denominator")
            elif isinstance(right, int):
                rhs *= right
                lhs = left
                continue
        if op == "*":
            if isinstance(left, int):
                rhs /= left
                lhs = right
                continue
            elif isinstance(right, int):
                rhs /= right
                lhs = left
                continue
    return rhs

def part2():
    left, right = monkeys["root"].split(" + ")

    print("Part 2:", int(solver(simplify(left), simplify(right))))

part1()
part2()