solve.py

import sys
from operator import add, mul, sub, truediv
from re import split

from sympy import solve as sympy_solve
from sympy.abc import x


def solve():
    monkeys = {}

    for line in sys.stdin:
        match split(r":? ", line.strip()):
            case [monkey, number]:
                monkeys[monkey] = int(number)
            case [monkey, arg1, operator, arg2]:
                monkeys[monkey] = [operator, arg1, arg2]

    print("Part 1:", int(resolve(monkeys, "root")))
    print("Part 2:", find(monkeys))


def resolve(monkeys, monkey):
    operations = {"+": add, "-": sub, "*": mul, "/": truediv}

    match monkeys[monkey]:
        case int(arg1):
            return arg1
        case [operator, arg1, arg2]:
            return operations[operator](resolve(monkeys, arg1), resolve(monkeys, arg2))


def find(monkeys):
    def build(monkey):
        match monkey:
            case "x":
                return x
            case int(monkey) | float(monkey):
                return monkey
            case [operator, *names]:
                left, right = [build(monkeys[name]) for name in names]

                return f"({left} {operator} {right})"

    monkeys["humn"] = "x"
    left, right = [build(monkeys[name]) for name in monkeys["root"][1:]]

    return sympy_solve(f"Eq({left}, {right})", x)[0]


if __name__ == "__main__":
    solve()