AOC 2023, day 14

aoc_2023_day14.py

# -------- #
# Part One #
# -------- #

def find_slices(col):
    slices = []
    fixed = np.argwhere(col == "#").ravel()
    for start, end in zip(fixed+1, fixed[1:]):
        if end > start:
            slices.append(slice(start, end))
    return slices

grid = np.pad(np.array(list(map(list, inp.splitlines()))), 1, constant_values="#")

for col in grid.T[1:]:
    for sl in find_slices(col):
        if n_roll := (col[sl] == "O").sum():
            col[sl][:n_roll] = "O"
            col[sl][n_roll:] = "."

# Solution, part 1:
np.sum((grid[1:-1, 1:-1] == "O").sum(axis=1) * np.arange(grid.shape[0]-2, 0, -1))


# -------- #
# Part Two #
# -------- #

import numpy as np
from numba import njit

def find_slices(col):
    slices = []
    fixed = np.argwhere(col == "#").ravel()
    for start, end in zip(fixed+1, fixed[1:]):
        if end > start:
            slices.append((start, end))
    return np.array(slices)

@njit
def let_roll_and_rock(grid, all_slices):
    for i, col in enumerate(grid.T[1:-1]):
        for start, end in all_slices[i]:
            sl = slice(start, end)
            n_roll = col[sl].sum()
            if n_roll > 0:
                col[sl][:n_roll] = 1
                col[sl][n_roll:] = 0
                
full_grid = np.pad(np.array(list(map(list, inp.splitlines()))), 1, constant_values="#")
rotated_grid = [full_grid, np.rot90(full_grid, 1, (1, 0)), np.rot90(full_grid, 2), np.rot90(full_grid, 1)]
all_slices = [tuple(find_slices(col) for col in grid.T[1:-1]) for grid in rotated_grid]
        
num_grid = np.zeros_like(full_grid, dtype=np.int8)
num_grid[full_grid == "O"] = 1
full_grid = num_grid
rotated_grid = (full_grid, np.rot90(full_grid, 1, (1, 0)), np.rot90(full_grid, 2), np.rot90(full_grid, 1))

stop_at, hashes = None, {}
for cycle in range(1, 1_000_000_000+1):
    for grid, slices in zip(rotated_grid, all_slices):
        let_roll_and_rock(grid, slices)
    
    if stop_at is None:
        # Only compute hash before first collision
        arr_hash = full_grid[1:-1, 1:-1].tobytes()
        if arr_hash in hashes:
            length = cycle - hashes[arr_hash]
            stop_at = cycle + (1_000_000_000 -  cycle) % length
        else:
            hashes[arr_hash] = cycle
    
    if stop_at == cycle:
        # Part 2:
        print(np.sum((full_grid[1:-1, 1:-1]).sum(axis=1) * np.arange(full_grid.shape[0]-2, 0, -1)))
        break