day19 slightly improved

advent/days/day19/solution.py

@@ -1,12 +1,13 @@
 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
@@ -71,51 +72,43 @@ 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 Path:
+class State:
     time: int
     material: Elements
     robots: Elements
     blueprint: Blueprint
-    path: list[Element | None]
 
     @classmethod
-    def start(cls, blueprint: Blueprint) -> Path:
+    def start(cls, blueprint: Blueprint) -> State:
-        return Path(0, (0, 0, 0, 0), (0, 0, 0, 1), blueprint, [])
+        return State(0, (0, 0, 0, 0), (0, 0, 0, 1), blueprint)
 
-    def _check(self, element: Element) -> Path | None:
+    def get_valid_production(self, element: Element) -> State | None:
         if element != Element.Geode:
-            max_needed = max(req[element] for req in self.blueprint.requirements)
+            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 Path(self.time + 1,
+            return State(self.time + 1,
-                        add(sub(self.material, self.blueprint.requirements[element]), self.robots),
+                         add(sub(self.material, self.blueprint.requirements[element]), self.robots),
-                        inc_element(self.robots, element), self.blueprint, self.path + [element])
+                         inc_element(self.robots, element), self.blueprint)
         else:
             return None
 
-    def find_next(self) -> Iterator[Path]:
+    def find_next(self) -> Iterator[State]:
-        if (path := self._check(Element.Geode)) is not None:
+        for element in Element:
-            yield path
+            if (path := self.get_valid_production(element)) is not None:
-
-        if self.blueprint.max_requirement(Element.Obsidian) >= self.material[Element.Obsidian]:
-            if (path := self._check(Element.Obsidian)) is not None:
                 yield path
 
-        if self.blueprint.max_requirement(Element.Clay) >= self.material[Element.Clay]:
+        yield State(self.time + 1, add(self.material, self.robots), self.robots, self.blueprint)
-            if (path := self._check(Element.Clay)) is not None:
-                yield path
 
-        if self.blueprint.max_requirement(Element.Ore) >= self.material[Element.Ore]:
+    def __lt__(self, other: State) -> bool:
-            if (path := self._check(Element.Ore)) is not None:
-                yield path
-
-        yield Path(self.time + 1, add(self.material, self.robots), self.robots, self.blueprint,
-                   self.path + [None])
-
-    def __lt__(self, other: Path) -> bool:
         if self.time != other.time:
             return self.time < other.time
         return self.material > other.material
@@ -125,6 +118,7 @@ class Path:
 class Blueprint:
     number: int
     requirements: tuple[Elements, Elements, Elements, Elements]
+    max_needed: Elements
 
     @classmethod
     def create(cls, number: int, ore: int, clay: int,
@@ -135,18 +129,16 @@ class Blueprint:
             (0, 0, 0, clay),
             (0, 0, 0, ore),
         )
-        return Blueprint(number, requirements)
+        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 max_requirement(self, element: Element) -> int:
-        return round(max(requirement[element] for requirement in self.requirements) * 1.2)
-
     def run(self, rounds: int) -> int:
-        queue: PriorityQueue[Path] = PriorityQueue()
+        queue: PriorityQueue[State] = PriorityQueue()
-        queue.put(Path.start(self))
+        queue.put(State.start(self))
         seen: dict[tuple[Elements, int], Elements] = {}
         while not queue.empty():
             current = queue.get()

advent/days/day19/test_solution.py

@@ -4,7 +4,6 @@ from .solution import Blueprint, day_num, part1, part2
 import pytest
 
 
-@pytest.mark.skip
 def test_part1():
     lines = input.read_lines(day_num, 'example01.txt')
     expected = 33
@@ -23,11 +22,11 @@ def test_part2():
 def test_parse():
     lines = input.read_lines(day_num, 'example01.txt')
     blueprint = Blueprint.parse(next(lines))
-    expected = Blueprint(1, ((0, 7, 0, 2), (0, 0, 14, 3), (0, 0, 0, 2), (0, 0, 0, 4)))
+    expected = Blueprint(1, ((0, 7, 0, 2), (0, 0, 14, 3),
+                         (0, 0, 0, 2), (0, 0, 0, 4)), (0, 7, 14, 4))
     assert blueprint == expected
 
 
-@pytest.mark.skip
 def test_blueprint1():
     lines = input.read_lines(day_num, 'example01.txt')
     blueprint = Blueprint.parse(next(lines))