aoc2021-24.py

import sys
import re

data = sys.stdin.read()

PARAMS = list(
    zip(
        map(int, re.findall("div z (26|1)", data)),
        map(int, re.findall("add x (-?[0-9]+)", data)),
        map(int, re.findall("add y w\nadd y (-?[0-9]+)", data, re.M)),
    )
)

for a, b, c in PARAMS:
    # always div by 26 or 1
    assert a in (1, 26)

    # b is always outside 0..10
    assert not (0 < b < 10)

    # a and b are connected
    assert (a == 1 and b >= 10) or (a == 26 and b <= 0)


def gen_relations():
    ret = {}
    stack = []
    for digit, (a, b, c) in enumerate(PARAMS):
        if a == 1:
            stack.append((digit, c))
        else:
            assert stack
            sd, sc = stack.pop()
            ret[digit] = sd, sc + b
            # print(f"# number[{digit}] = number[{ret[digit][0]}] + {ret[digit][0]} + {b}")
    assert not stack
    return ret


def solve(smallest=False):
    rules = gen_relations()

    DIGIT_ORDER = list(range(9, 0, -1))
    if smallest:
        DIGIT_ORDER = list(reversed(DIGIT_ORDER))

    def solve_recursive(digits={}):
        if len(digits) >= 14:
            return digits

        digit = min(set(range(14)) - set(digits))

        if digit in rules:
            rule = rules[digit]
            value = digits[rule[0]] + rule[1]
            if 0 < value <= 9:
                digits[digit] = value
                solution = solve_recursive(digits)
                if solution:
                    return solution
        else:
            for value in DIGIT_ORDER:
                digits[digit] = value
                solution = solve_recursive(digits)
                if solution:
                    return solution

        digits.pop(digit, None)

    solution = solve_recursive()
    if solution:
        solution = int("".join(str(solution[digit]) for digit in range(14)))
        return solution


print("Part 1:", solve())
print("Part 2:", solve(True))