2024-12-09 10:33:28 +01:00

68 lines
2.2 KiB
Python

from typing import Any, Iterator
from ..base import BaseSolver
class Solver(BaseSolver):
def solve(self, input: str) -> Iterator[Any]:
blocks: list[tuple[int, int]] = []
frees: list[tuple[int, int]] = []
contents_0: list[int | None] = [None for _ in range(sum(map(int, input)))]
acc = 0
for i, c in enumerate(input):
if i % 2 == 0:
for j in range(acc, acc + int(c)):
contents_0[j] = i // 2
blocks.append((acc, int(c)))
else:
frees.append((acc, int(c)))
acc += int(c)
assert contents_0[-1] is not None
contents = contents_0.copy()
free_0 = next(i for i, c in enumerate(contents) if c is None)
next_b = len(contents) - 1
while free_0 < next_b:
contents[free_0], contents[next_b] = contents[next_b], contents[free_0]
free_0 += 1
while free_0 < len(contents) and contents[free_0] is not None:
free_0 += 1
next_b -= 1
while next_b >= 0 and contents[next_b] is None:
next_b -= 1
yield sum(i * c for i, c in enumerate(contents) if c is not None)
contents = contents_0.copy()
for block_start, block_length in self.progress.wrap(blocks[::-1]):
try:
i_free = next(
i_free
for i_free, (free_start, free_length) in enumerate(frees)
if free_start < block_start and free_length >= block_length
)
except StopIteration:
continue
free_start, free_length = frees[i_free]
contents[free_start : free_start + block_length] = contents[
block_start : block_start + block_length
]
contents[block_start : block_start + block_length] = [None] * block_length
if free_length == block_length:
del frees[i_free]
else:
frees[i_free] = (free_start + block_length, free_length - block_length)
yield sum(i * c for i, c in enumerate(contents) if c is not None)