121 lines
3.3 KiB
Python
121 lines
3.3 KiB
Python
import sys
|
|
from typing import Sequence, TypeVar
|
|
|
|
import numpy as np
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def print_tower(tower: np.ndarray, out: str = "#"):
|
|
print("-" * (tower.shape[1] + 2))
|
|
non_empty = False
|
|
for row in reversed(range(1, tower.shape[0])):
|
|
if not non_empty and not tower[row, :].any():
|
|
continue
|
|
non_empty = True
|
|
print("|" + "".join(out if c else "." for c in tower[row, :]) + "|")
|
|
print("+" + "-" * tower.shape[1] + "+")
|
|
|
|
|
|
def tower_height(tower: np.ndarray) -> int:
|
|
return int(tower.shape[0] - tower[::-1, :].argmax(axis=0).min() - 1)
|
|
|
|
|
|
def next_cycle(sequence: Sequence[T], index: int) -> tuple[T, int]:
|
|
t = sequence[index]
|
|
index = (index + 1) % len(sequence)
|
|
return t, index
|
|
|
|
|
|
ROCKS = [
|
|
np.array([(0, 0), (0, 1), (0, 2), (0, 3)]),
|
|
np.array([(0, 1), (1, 0), (1, 1), (1, 2), (2, 1)]),
|
|
np.array([(0, 0), (0, 1), (0, 2), (1, 2), (2, 2)]),
|
|
np.array([(0, 0), (1, 0), (2, 0), (3, 0)]),
|
|
np.array([(0, 0), (0, 1), (1, 0), (1, 1)]),
|
|
]
|
|
|
|
WIDTH = 7
|
|
START_X = 2
|
|
|
|
EMPTY_BLOCKS = np.zeros((10, WIDTH), dtype=bool)
|
|
|
|
|
|
def build_tower(
|
|
n_rocks: int,
|
|
jets: str,
|
|
early_stop: bool = False,
|
|
init: np.ndarray = np.ones(WIDTH, dtype=bool),
|
|
) -> tuple[np.ndarray, int, int, dict[int, int]]:
|
|
tower = EMPTY_BLOCKS.copy()
|
|
tower[0, :] = init
|
|
|
|
done_at: dict[tuple[int, int], int] = {}
|
|
heights: dict[int, int] = {}
|
|
i_jet, i_rock = 0, 0
|
|
rock_count = 0
|
|
|
|
for rock_count in range(n_rocks):
|
|
if early_stop:
|
|
if i_rock == 0 and (i_rock, i_jet) in done_at:
|
|
break
|
|
done_at[i_rock, i_jet] = rock_count
|
|
|
|
y_start = tower.shape[0] - tower[::-1, :].argmax(axis=0).min() + 3
|
|
rock, i_rock = next_cycle(ROCKS, i_rock)
|
|
|
|
rock_y = rock[:, 0] + y_start
|
|
rock_x = rock[:, 1] + START_X
|
|
|
|
if rock_y.max() >= tower.shape[0]:
|
|
tower = np.concatenate([tower, EMPTY_BLOCKS], axis=0)
|
|
|
|
while True:
|
|
jet, i_jet = next_cycle(jets, i_jet)
|
|
|
|
dx = 0
|
|
if jet == ">" and rock_x.max() < WIDTH - 1:
|
|
dx = 1
|
|
elif jet == "<" and rock_x.min() > 0:
|
|
dx = -1
|
|
|
|
if dx != 0 and not tower[rock_y, rock_x + dx].any():
|
|
rock_x = rock_x + dx
|
|
|
|
# move down
|
|
rock_y -= 1
|
|
|
|
if tower[rock_y, rock_x].any():
|
|
rock_y += 1
|
|
break
|
|
|
|
heights[rock_count] = tower_height(tower)
|
|
tower[rock_y, rock_x] = True
|
|
|
|
return tower, rock_count, done_at.get((i_rock, i_jet), -1), heights
|
|
|
|
|
|
line = sys.stdin.read().strip()
|
|
|
|
tower, *_ = build_tower(2022, line)
|
|
answer_1 = tower_height(tower)
|
|
print(f"answer 1 is {answer_1}")
|
|
|
|
TOTAL_ROCKS = 1_000_000_000_000
|
|
tower_1, n_rocks_1, prev_1, heights_1 = build_tower(TOTAL_ROCKS, line, True)
|
|
assert prev_1 > 0
|
|
|
|
# 2767 1513
|
|
remaining_rocks = TOTAL_ROCKS - n_rocks_1
|
|
n_repeat_rocks = n_rocks_1 - prev_1
|
|
n_repeat_towers = remaining_rocks // n_repeat_rocks
|
|
|
|
base_height = heights_1[prev_1]
|
|
repeat_height = heights_1[prev_1 + n_repeat_rocks - 1] - heights_1[prev_1]
|
|
remaining_height = (
|
|
heights_1[prev_1 + remaining_rocks % n_repeat_rocks] - heights_1[prev_1]
|
|
)
|
|
|
|
answer_2 = base_height + (n_repeat_towers + 1) * repeat_height + remaining_height
|
|
print(f"answer 2 is {answer_2}")
|