advent-of-code/2023/day19.py

import logging
import operator
import os
import sys
from math import prod
from typing import Literal, TypeAlias, cast

VERBOSE = os.getenv("AOC_VERBOSE") == "True"

logging.basicConfig(level=logging.INFO if VERBOSE else logging.WARNING)

Category: TypeAlias = Literal["x", "m", "a", "s"]
Part: TypeAlias = dict[Category, int]
PartWithBounds: TypeAlias = dict[Category, tuple[int, int]]

OPERATORS = {"<": operator.lt, ">": operator.gt}

# None if there is no check (last entry), otherwise (category, sense, value)
Check: TypeAlias = tuple[Category, Literal["<", ">"], int] | None

# workflow as a list of check, in specified order, with target
Workflow: TypeAlias = list[tuple[Check, str]]


def accept(workflows: dict[str, Workflow], part: Part) -> bool:
    workflow = "in"
    decision: bool | None = None

    while decision is None:
        for check, target in workflows[workflow]:
            ok = check is None
            if check is not None:
                category, sense, value = check
                ok = OPERATORS[sense](part[category], value)

            if ok:
                if target in workflows:
                    workflow = target
                else:
                    decision = target == "A"
                break

    return decision


def propagate(workflows: dict[str, Workflow], start: PartWithBounds) -> int:
    def transfer_or_accept(
        target: str, meta: PartWithBounds, queue: list[tuple[PartWithBounds, str]]
    ) -> int:
        if target in workflows:
            logging.info(f"  transfer to {target}")
            queue.append((meta, target))
            return 0
        elif target == "A":
            logging.info("  accepted")
            return prod((high - low + 1) for low, high in meta.values())
        else:
            logging.info("  rejected")
            return 0

    accepted = 0
    queue: list[tuple[PartWithBounds, str]] = [(start, "in")]

    while queue:
        meta, workflow = queue.pop()
        logging.info(f"{workflow}: {meta}")
        for check, target in workflows[workflow]:
            if check is None:
                accepted += transfer_or_accept(target, meta, queue)
                continue

            category, sense, value = check
            bounds, op = meta[category], OPERATORS[sense]

            logging.info(f"  splitting {meta} into {category} {sense} {value}")

            if not op(bounds[0], value) and not op(bounds[1], value):
                logging.info("  reject, always false")
                continue

            if op(meta[category][0], value) and op(meta[category][1], value):
                logging.info("  accept, always true")
                accepted += transfer_or_accept(target, meta, queue)
                break

            meta2 = meta.copy()
            if sense == "<":
                meta2[category] = (meta[category][0], value - 1)
                meta[category] = (value, meta[category][1])
            else:
                meta2[category] = (value + 1, meta[category][1])
                meta[category] = (meta[category][0], value)
            logging.info(f"  split {meta2} ({target}), {meta}")

            accepted += transfer_or_accept(target, meta2, queue)

    return accepted


workflows_s, parts_s = sys.stdin.read().strip().split("\n\n")

workflows: dict[str, Workflow] = {}
for workflow_s in workflows_s.split("\n"):
    name, block_s = workflow_s.split("{")
    workflows[name] = []

    for block in block_s[:-1].split(","):
        check: Check
        if (i := block.find(":")) >= 0:
            check, target = (
                cast(Category, block[0]),
                cast(Literal["<", ">"], block[1]),
                int(block[2:i]),
            ), block[i + 1 :]
        else:
            check, target = None, block
        workflows[name].append((check, target))

# part 1
parts: list[Part] = [
    {cast(Category, s[0]): int(s[2:]) for s in part_s[1:-1].split(",")}
    for part_s in parts_s.split("\n")
]
answer_1 = sum(sum(part.values()) for part in parts if accept(workflows, part))
print(f"answer 1 is {answer_1}")


# part 2
answer_2 = propagate(
    workflows, {cast(Category, c): (1, 4000) for c in ["x", "m", "a", "s"]}
)
print(f"answer 2 is {answer_2}")