2022: d19: ex2: add solution

2022/d19/ex2/ex2.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python
+
+import dataclasses
+import enum
+import itertools
+import math
+import sys
+from collections import deque
+from collections.abc import Iterable, Iterator, Mapping
+from typing import NamedTuple, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+def grouper(iterable: Iterable[T], n: int) -> Iterator[tuple[T, ...]]:
+    "Collect data into non-overlapping fixed-length chunks or blocks"
+    args = [iter(iterable)] * n
+    return zip(*args, strict=True)
+
+
+class Resource(str, enum.Enum):
+    GEODE = "geode"
+    OBSIDIAN = "obsidian"
+    CLAY = "clay"
+    ORE = "ore"
+
+
+class ResourceCost(Mapping[Resource, int]):
+    _dict: dict[Resource, int]
+    _hash: Optional[int]
+
+    def __init__(self, init: Mapping[Resource, int] = {}, /) -> None:
+        self._dict = {res: init.get(res, 0) for res in Resource}
+        self._hash = None
+
+        assert all(self._dict[res] >= 0 for res in Resource)  # Sanity check
+
+    def __getitem__(self, key: Resource, /) -> int:
+        return self._dict[key]
+
+    def __iter__(self) -> Iterator[Resource]:
+        return iter(Resource)  # Always use same Resource iteration order
+
+    def __len__(self) -> int:
+        return len(self._dict)
+
+    def __hash__(self) -> int:
+        if self._hash is None:
+            self._hash = hash(tuple(sorted(self._dict)))
+        return self._hash
+
+    def __add__(self, other):
+        if not isinstance(other, ResourceCost):
+            return NotImplemented
+        return ResourceCost({res: self[res] + other[res] for res in Resource})
+
+    def __sub__(self, other):
+        if not isinstance(other, ResourceCost):
+            return NotImplemented
+        return ResourceCost({res: self[res] - other[res] for res in Resource})
+
+    def __repr__(self) -> str:
+        return repr(self._dict)
+
+    def has_enough(self, costs: "ResourceCost") -> bool:
+        return all(self[res] >= costs[res] for res in Resource)
+
+
+@dataclasses.dataclass
+class Blueprint:
+    construction_costs: dict[Resource, ResourceCost]
+
+    @classmethod
+    def from_input(cls, input: str) -> "Blueprint":
+        assert input.startswith("Blueprint ")  # Sanity check
+
+        raw_costs = input.split(": ")[1].split(". ")
+        costs: dict[Resource, ResourceCost] = {}
+        for raw in map(str.split, raw_costs):
+            ressource = Resource(raw[1])
+            costs[ressource] = ResourceCost(
+                {
+                    Resource(r.removesuffix(".")): int(c)
+                    for c, r in grouper((w for w in raw[4:] if w != "and"), 2)
+                }
+            )
+
+        return cls(costs)
+
+    def maximize_geodes(self, run_time: int) -> int:
+        class QueueNode(NamedTuple):
+            time: int
+            robots: ResourceCost
+            inventory: ResourceCost
+            total_mined: ResourceCost
+
+        def prune_queue(queue: Iterable[QueueNode]) -> deque[QueueNode]:
+            def priority_key(node: QueueNode) -> int:
+                MULTIPLIERS = {
+                    Resource.GEODE: 1_000_000,
+                    Resource.OBSIDIAN: 10_000,
+                    Resource.CLAY: 100,
+                    Resource.ORE: 1,
+                }
+                return sum(
+                    node.total_mined[res] * mul for res, mul in MULTIPLIERS.items()
+                )
+
+            MAX_QUEUE = 10_000  # Chosen arbitrarily
+            return deque(sorted(queue, key=priority_key, reverse=True)[:MAX_QUEUE])
+
+        def do_build(node: QueueNode, robot_type: Optional[Resource]) -> QueueNode:
+            costs = (
+                self.construction_costs[robot_type]
+                if robot_type is not None
+                else ResourceCost()
+            )
+            assert node.inventory.has_enough(costs)  # Sanity check
+            new_robots = node.robots + (
+                ResourceCost({robot_type: 1})
+                if robot_type is not None
+                else ResourceCost()
+            )
+            new_inventory = node.inventory + node.robots - costs
+            new_total_mined = node.total_mined + node.robots
+            return QueueNode(node.time + 1, new_robots, new_inventory, new_total_mined)
+
+        max_geode = 0
+
+        queue: deque[QueueNode] = deque(
+            # Starting conditions
+            [
+                QueueNode(
+                    0, ResourceCost({Resource.ORE: 1}), ResourceCost(), ResourceCost()
+                )
+            ]
+        )
+        dfs_depth = 0
+        while queue:
+            node = queue.popleft()
+
+            if node.time > dfs_depth:
+                # An awful hack to reduce the search space and prioritize geodes
+                queue = prune_queue(queue)
+                dfs_depth = node.time
+
+            if node.time == run_time:
+                max_geode = max(max_geode, node.total_mined[Resource.GEODE])
+                continue
+
+            # Try building a robot
+            for robot_type in itertools.chain(Resource):
+                costs = self.construction_costs[robot_type]
+                # Don't build robots we can't afford
+                if not node.inventory.has_enough(costs):
+                    continue
+                # Don't build robots when already producing more than enough
+                if robot_type != Resource.GEODE and all(
+                    c[robot_type] <= node.robots[robot_type]
+                    for c in self.construction_costs.values()
+                ):
+                    continue
+                queue.append(do_build(node, robot_type))
+            # Try not building anything
+            queue.append(do_build(node, None))
+
+        return max_geode
+
+
+def solve(input: list[str]) -> int:
+    blueprints = [Blueprint.from_input(line) for line in input]
+
+    TIME = 32
+
+    return math.prod(blueprint.maximize_geodes(TIME) for blueprint in blueprints[:3])
+
+
+def main() -> None:
+    input = sys.stdin.read().splitlines()
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()