187 lines
5.8 KiB
Python
Executable file
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()
|