155 lines
5 KiB
Python
155 lines
5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import IntEnum
|
|
from itertools import islice
|
|
from math import prod
|
|
from multiprocessing import Pool
|
|
from queue import PriorityQueue
|
|
from typing import Iterable, Iterator, Self
|
|
|
|
from advent.parser.parser import P
|
|
|
|
day_num = 19
|
|
|
|
|
|
def part1(lines: Iterator[str]) -> int:
|
|
return sum(bp.number * geodes for bp, geodes in Processor.pool_it(24, lines))
|
|
|
|
|
|
def part2(lines: Iterator[str]) -> int:
|
|
return prod(num for _, num in Processor.pool_it(32, islice(lines, 3)))
|
|
|
|
|
|
@dataclass(slots=True, frozen=True)
|
|
class Processor:
|
|
rounds: int
|
|
|
|
@classmethod
|
|
def pool_it(cls, rounds: int, lines: Iterator[str]) -> Iterable[tuple[Blueprint, int]]:
|
|
with Pool() as p:
|
|
return p.map(Processor(rounds), lines)
|
|
|
|
def __call__(self, line: str) -> tuple[Blueprint, int]:
|
|
blueprint = Blueprint.parse(line)
|
|
return blueprint, blueprint.run(self.rounds)
|
|
|
|
|
|
number_parser = P.second(P.string("Blueprint "), P.unsigned())
|
|
ore_parser = P.second(P.string(": Each ore robot costs "), P.unsigned())
|
|
clay_parser = P.second(P.string(" ore. Each clay robot costs "), P.unsigned())
|
|
tuple_parser = P.seq(P.unsigned(), P.second(P.string(" ore and "), P.unsigned()))
|
|
obsidian_parser = P.second(P.string(" ore. Each obsidian robot costs "), tuple_parser)
|
|
geode_parser = P.second(P.string(" clay. Each geode robot costs "), tuple_parser)
|
|
blueprint_parser = P.map5(number_parser, ore_parser, clay_parser, obsidian_parser, geode_parser,
|
|
lambda number, ore, clay, obsidian, geode:
|
|
Blueprint.create(number, ore, clay, obsidian, geode))
|
|
|
|
|
|
class Element(IntEnum):
|
|
Geode = 0
|
|
Obsidian = 1
|
|
Clay = 2
|
|
Ore = 3
|
|
|
|
|
|
Elements = tuple[int, int, int, int]
|
|
|
|
|
|
def gt(first: Elements, second: Elements) -> bool:
|
|
return all(f >= s for f, s in zip(first, second))
|
|
|
|
|
|
def add(first: Elements, second: Elements) -> Elements:
|
|
return tuple(v1 + v2 for v1, v2 in zip(first, second))
|
|
|
|
|
|
def sub(first: Elements, second: Elements) -> Elements:
|
|
return Elements(tuple(v1 - v2 for v1, v2 in zip(first, second)))
|
|
|
|
|
|
def inc_element(elements: Elements, pos: Element) -> Elements:
|
|
return tuple(v + 1 if num == pos else v for num, v in enumerate(elements))
|
|
|
|
|
|
MAGIC_MATERIAL_SURPLUS = 1.2
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class State:
|
|
time: int
|
|
material: Elements
|
|
robots: Elements
|
|
blueprint: Blueprint
|
|
|
|
@classmethod
|
|
def start(cls, blueprint: Blueprint) -> State:
|
|
return State(0, (0, 0, 0, 0), (0, 0, 0, 1), blueprint)
|
|
|
|
def get_valid_production(self, element: Element) -> State | None:
|
|
if element != Element.Geode:
|
|
max_needed = self.blueprint.max_needed[element]
|
|
if self.robots[element] >= max_needed:
|
|
return None
|
|
if self.material[element] > round(max_needed * MAGIC_MATERIAL_SURPLUS):
|
|
return None
|
|
|
|
if gt(self.material, self.blueprint.requirements[element]):
|
|
return State(self.time + 1,
|
|
add(sub(self.material, self.blueprint.requirements[element]), self.robots),
|
|
inc_element(self.robots, element), self.blueprint)
|
|
else:
|
|
return None
|
|
|
|
def find_next(self) -> Iterator[State]:
|
|
for element in Element:
|
|
if (path := self.get_valid_production(element)) is not None:
|
|
yield path
|
|
|
|
yield State(self.time + 1, add(self.material, self.robots), self.robots, self.blueprint)
|
|
|
|
def __lt__(self, other: State) -> bool:
|
|
if self.time != other.time:
|
|
return self.time < other.time
|
|
return self.material > other.material
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class Blueprint:
|
|
number: int
|
|
requirements: tuple[Elements, Elements, Elements, Elements]
|
|
max_needed: Elements
|
|
|
|
@classmethod
|
|
def create(cls, number: int, ore: int, clay: int,
|
|
obsidian: tuple[int, int], geode: tuple[int, int]) -> Self:
|
|
requirements = (
|
|
(0, geode[1], 0, geode[0]),
|
|
(0, 0, obsidian[1], obsidian[0]),
|
|
(0, 0, 0, clay),
|
|
(0, 0, 0, ore),
|
|
)
|
|
max_needed = (0, geode[1], obsidian[1], max(geode[0], obsidian[0], clay, ore))
|
|
return Blueprint(number, requirements, max_needed)
|
|
|
|
@classmethod
|
|
def parse(cls, line: str) -> Self:
|
|
return blueprint_parser.parse(line).get()
|
|
|
|
def run(self, rounds: int) -> int:
|
|
queue: PriorityQueue[State] = PriorityQueue()
|
|
queue.put(State.start(self))
|
|
seen: dict[tuple[Elements, int], Elements] = {}
|
|
while not queue.empty():
|
|
current = queue.get()
|
|
last_seen = seen.get((current.robots, current.time))
|
|
if last_seen is not None and gt(last_seen, current.material):
|
|
continue
|
|
seen[(current.robots, current.time)] = current.material
|
|
|
|
if current.time == rounds:
|
|
return current.material[Element.Geode]
|
|
for next in current.find_next():
|
|
queue.put(next)
|
|
|
|
raise Exception("No optimum found")
|