diff --git a/2022/d19/ex1/ex1.py b/2022/d19/ex1/ex1.py new file mode 100755 index 0000000..9b663a7 --- /dev/null +++ b/2022/d19/ex1/ex1.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python + +import dataclasses +import enum +import itertools +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 = 1000 # 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 = 24 + + return sum( + i * blueprint.maximize_geodes(TIME) + for i, blueprint in enumerate(blueprints, start=1) + ) + + +def main() -> None: + input = sys.stdin.read().splitlines() + print(solve(input)) + + +if __name__ == "__main__": + main()