advent-of-code/2022/d19/ex1/ex1.py

187 lines
5.8 KiB
Python
Executable file

#!/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()