Re: Optimisation

Avec ladite optimisation :

from __future__ import annotations
import enum
import io
import itertools
from typing import Iterable, Iterator, List, Tuple
import numpy as np
import numpy.typing as npt
Coords = Tuple[int, int]
class Material(enum.Enum):
 AIR = enum.auto()
 ROCK = enum.auto()
 SAND = enum.auto()
class Slice:
 def __init__(self, array: npt.ArrayLike, leak: Coords) -> None:
 self.matrix: npt.NDArray = np.array(array)
 self.leak = leak
 def pour(self) -> Iterator[bool]:
 """Inject a unit of sand.
 Yields False if it was blocked, True if it fell into the void.
 Stops when the leak is clogged up by a mountain of sand."""
 # List of coordinates to pour sand from
 sources: List[Coords] = []
 # Start at the original leak
 y, x = self.leak
 while True:
 # Our unit of sand is injected at current coordinates (y, x)
 while y < self.matrix.shape[0] - 1:
 y_ = y + 1
 for x_ in (x, x - 1, x + 1):
 if self.matrix[y_, x_] is Material.AIR:
 # Sand can flow down:
 # * current coordinate is a good place to pour next
 # sand unit from;
 sources.append((y, x))
 # * current sand unit falls down one level.
 y = y_
 x = x_
 break
 else:
 # Ways down exhausted, sand is blocked
 self.matrix[y, x] = Material.SAND
 # Get out of the descent loop
 break
 else:
 # Descent ended, sand fell all the way down
 yield True
 continue
 # Descent did not end, sand got blocked
 yield False
 if not sources:
 # Nowhere to pour sand from, even the leak is clogged up
 break
 # Pour next sand unit from last good place
 y, x = sources.pop()
 def __str__(self) -> str:
 result = io.StringIO()
 for line in self.matrix:
 for point in line:
 if point is Material.AIR:
 result.write(' ')
 elif point is Material.ROCK:
 result.write('█')
 elif point is Material.SAND:
 result.write('░')
 else:
 assert False # we covered all cases
 result.write('\n')
 return result.getvalue()
 def add_segment(self, start: Coords, end: Coords) -> None:
 y1, x1 = start
 y2, x2 = end
 if y1 == y2:
 x1, x2 = min(x1, x2), max(x1, x2)
 ys = (y1 for _ in range(x1, x2 + 1))
 xs = (x for x in range(x1, x2 + 1))
 elif x1 == x2:
 y1, y2 = min(y1, y2), max(y1, y2)
 ys = (y for y in range(y1, y2 + 1))
 xs = (x1 for _ in range(y1, y2 + 1))
 for y, x in zip(ys, xs):
 self.matrix[y, x] = Material.ROCK
def import_segments(
 lines: Iterable[str]
 ) -> Tuple[List[Coords], List[Tuple[Coords, Coords]]]:
 points: List[Coords] = []
 segments: List[Tuple[Coords, Coords]] = []
 def coords(word: str) -> Coords:
 parts = word.split(',')
 if len(parts) != 2:
 raise ValueError("unexpected number of coordinates")
 return int(parts[1]), int(parts[0])
 for line in lines:
 words = line.rstrip().split(' -> ')
 segment_points = [coords(word) for word in words]
 points.extend(segment_points)
 segments.extend(itertools.pairwise(segment_points))
 return points, segments
def import_slice1(lines: Iterable[str]) -> Slice:
 points, segments = import_segments(lines)
 y_leak, x_leak = 0, 500
 points.append((y_leak, x_leak))
 xmin = min(x for _, x in points) - 1
 xmax = max(x for _, x in points) + 1
 ymin = min(y for y, _ in points)
 ymax = max(y for y, _ in points)
 if ymin < 0:
 raise ValueError("unexpected negative coordinate")
 xshift = xmin
 matrix = np.full(((ymax + 1), xmax - xshift + 1), Material.AIR)
 result = Slice(matrix, (0, 500 - xshift))
 for (y1, x1), (y2, x2) in segments:
 result.add_segment((y1, x1 - xshift), (y2, x2 - xshift))
 return result
def import_slice2(lines: Iterable[str]) -> Slice:
 points, segments = import_segments(lines)
 y_leak, x_leak = 0, 500
 points.append((y_leak, x_leak))
 xmin = min(x for _, x in points)
 xmax = max(x for _, x in points)
 ymin = min(y for y, _ in points)
 ymax = max(y for y, _ in points) + 2 # including the floor
 # Update xmin and xmax to accomodate a pyramid of sand
 xmin = min(xmin, x_leak - (ymax - y_leak))
 xmax = max(xmax, x_leak + (ymax - y_leak))
 if ymin < 0:
 raise ValueError("unexpected negative coordinate")
 xshift = xmin
 matrix = np.full(((ymax + 1), xmax - xshift + 1), Material.AIR)
 result = Slice(matrix, (0, 500 - xshift))
 for (y1, x1), (y2, x2) in segments:
 result.add_segment((y1, x1 - xshift), (y2, x2 - xshift))
 result.add_segment((ymax, xmin - xshift), (ymax, xmax - xshift))
 return result
def solve1(lines: Iterable[str]) -> int:
 """Solve part 1 of today's puzzle"""
 slice_ = import_slice1(lines)
 for i, fallen in enumerate(slice_.pour()):
 if fallen:
 return i
 raise ValueError("Simulation ended with unexpected sand leak clogged")
def solve2(lines: Iterable[str]) -> int:
 """Solve part 2 of today's puzzle"""
 slice_ = import_slice2(lines)
 for i, fallen in enumerate(slice_.pour()):
 if fallen:
 raise ValueError("Simulation ended with unexpected sand under the floor")
 return i + 1