aoc-2024/7/1.py

from itertools import product

def test(expression: list[str|int]) -> int:
    while len(expression) >= 3:
        left, op, right = expression[0], expression[1], expression[2]
        if len(expression) > 3:
            expression = expression[3:]
        else:
            expression = []
        if op == '*':
            expression.insert(0, left * right)
        else:
            expression.insert(0, left + right)
    return expression[0]

def genrate_expression(vals: list[int], ops: list[str]) -> list[str|int]:
    expr = []
    for item in zip(vals, ops):
        expr.extend(item)
    expr.append(vals[-1])
    return expr


with open('sample.data') as file:
    inputs = []
    for line in file:
        vals = [int(val) for val in line.replace(':', '').split()]
        inputs.append([vals[0], vals[1:]])
    s = 0
    for case in inputs:
        answer, vals = case[0], case[1]
        for ops in list(product('*+', repeat=len(vals)-1)):
            expr = genrate_expression(vals, ops)
            result = test(expr)
            if result == answer:
                s += answer
                break
    print(s)