advent-of-code/src/holt59/aoc/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]:
            passed = check is None
            if check is not None:
                category, sense, value = check
                passed = OPERATORS[sense](part[category], value)

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

    return decision


def propagate(workflows: dict[str, Workflow], start: PartWithBounds) -> int:
    def _fmt(meta: PartWithBounds) -> str:
        return "{" + ", ".join(f"{k}={v}" for k, v in meta.items()) + "}"

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

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

    n_iterations = 0

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

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

            logging.info(f"  checking {_fmt(meta)} against {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()
            low, high = meta[category]
            if sense == "<":
                meta[category], meta2[category] = (value, high), (low, value - 1)
            else:
                meta[category], meta2[category] = (low, value), (value + 1, high)
            logging.info(f"    split {_fmt(meta2)} ({target}), {_fmt(meta)}")

            accepted += transfer_or_accept(target, meta2, queue)

    logging.info(f"run took {n_iterations} iterations")
    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 = (
                cast(Category, block[0]),
                cast(Literal["<", ">"], block[1]),
                int(block[2:i]),
            )
            target = 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}")