aoc2024/d07/run.py

import sys
from typing import Generator as Gen

PART2 = len(sys.argv)>1 and sys.argv[1] == "2"
I = sys.stdin.read().strip().split("\n")

def conccat(a, b) -> int:
    # 12 345 become 12345
    # return a * pow(10, ceil(log10(b))) + b <= breaksfor 1000
    # but int <-> str convertions are slow (x4)
    return int(str(a)+str(b))

def gen_nb(L:list[int]) -> Gen[int, None, None]:
    if len(L) == 1:
        yield L[0]
    else:
        for x in gen_nb(L[:-1]):
            yield x * L[-1]
            yield x + L[-1]
            if PART2:
                yield conccat(x, L[-1])

S = 0
for l in I:
    x, nb = l.split(":")
    x = int(x)
    nb = list(map(int,nb.strip().split(" ")))
    S += x * (1 if any(x == a for a in gen_nb(nb)) else 0)
print(S)