advent-of-code/src/holt59/aoc/2024/day7.py

from typing import Any, Iterator

from ..base import BaseSolver


def evaluate(
    target: int, numbers: list[int], concatenate: bool = False, current: int = 0
) -> bool:
    if not numbers:
        return current == target

    if current > target:
        return False

    head, *tail = numbers

    if evaluate(target, tail, concatenate, current + head) or evaluate(
        target, tail, concatenate, current * head
    ):
        return True

    if not concatenate:
        return False

    return evaluate(target, tail, concatenate, int(str(current) + str(head)))


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        targets = {
            int(part[0]): list(map(int, part[1].strip().split()))
            for line in input.splitlines()
            if (part := line.split(":"))
        }

        yield sum(
            target
            for target, numbers in self.progress.wrap(
                targets.items(), total=len(targets)
            )
            if evaluate(target, numbers)
        )

        yield sum(
            target
            for target, numbers in self.progress.wrap(
                targets.items(), total=len(targets)
            )
            if evaluate(target, numbers, True)
        )