adventofcode2024/day07/sol.py

def solve(file):
    t = 0
    t2 = 0
    for line  in file:
        nums = line.split()
        goal = int(nums[0].rstrip(":"))
        nums = [int(x) for x in nums[1:]]

        if viable(goal, nums, lambda n,x: (n*x, n+x)):
            t += goal
        elif viable(goal, nums, lambda n,x: (n*x, n+x, int(str(n)+str(x)))):
            t2 += goal
    print(t)
    print(t+t2)

def viable(goal, nums, combine):
    candidates = [nums[0]]
    for x in nums[1:-1]:
        candidates = (lambda C, x: (
            m
            for n in C
            for m in combine(n,x)
            if m < goal
        ))(candidates, x)
    #candidates = list(candidates); print(len(candidates))
    x = nums[-1]
    for n in candidates:
        for m in combine(n,x):
            if m == goal:
                return True
    return False

solve(open('sample1.in'))
solve(open('input'))