|
| 1 | +from itertools import product, cycle |
| 2 | + |
| 3 | +import numpy as np |
| 4 | +from scipy.ndimage import convolve |
| 5 | + |
| 6 | +def check(grid, cols, rows): |
| 7 | + for c, r in zip(cols, rows): |
| 8 | + k = grid[c, r] |
| 9 | + if k in (1,2): |
| 10 | + return k-1 |
| 11 | + return 0 |
| 12 | + |
| 13 | +def get_neightbors(grid, orig_grid, is_part1): |
| 14 | + if is_part1: |
| 15 | + kernel = np.ones((3,3), dtype=np.int8) |
| 16 | + kernel[1,1] = 0 |
| 17 | + empty = convolve(orig_grid, kernel, mode="constant", cval=0) |
| 18 | + return convolve(grid, kernel, mode="constant", cval=0) - empty |
| 19 | + else: |
| 20 | + res = np.zeros_like(grid) |
| 21 | + col, row = grid.shape |
| 22 | + for r, c in product(range(row), range(col)): |
| 23 | + if grid[c,r] == 0: |
| 24 | + continue |
| 25 | + else: |
| 26 | + n = check(grid, cycle([c]), reversed(range(r))) # left |
| 27 | + n += check(grid, cycle([c]), range(r+1, row)) # right |
| 28 | + n += check(grid, reversed(range(c)), cycle([r])) # up |
| 29 | + n += check(grid, range(c+1, col), cycle([r])) # down |
| 30 | + |
| 31 | + n += check(grid, reversed(range(c)), reversed(range(r))) # diag left up |
| 32 | + n += check(grid, range(c+1, col), reversed(range(r))) # diag left down |
| 33 | + n += check(grid, reversed(range(c)), range(r+1, row)) # diag right up |
| 34 | + n += check(grid, range(c+1, col), range(r+1, row)) # diag right down |
| 35 | + |
| 36 | + res[c,r] = n |
| 37 | + |
| 38 | + return res |
| 39 | + |
| 40 | +with open("input.txt") as f: |
| 41 | + rows = [] |
| 42 | + for x in f: |
| 43 | + rows.append([c == "L" for c in x.strip()]) |
| 44 | + |
| 45 | +for is_part1 in [True, False]: |
| 46 | + |
| 47 | + if is_part1: |
| 48 | + cutoff = 4 |
| 49 | + else: |
| 50 | + cutoff = 5 |
| 51 | + |
| 52 | + grid = np.array(rows).astype(np.int8) |
| 53 | + seats = grid.copy() |
| 54 | + occ = (grid == 2).sum() |
| 55 | + |
| 56 | + while True: |
| 57 | + prev_occ = occ |
| 58 | + n = get_neightbors(grid, seats, is_part1) |
| 59 | + u = (n == 0).astype(int) - (n >= cutoff).astype(int) |
| 60 | + u[grid == 0] = 0 |
| 61 | + grid = np.clip(grid + u, 0, 2) |
| 62 | + grid[grid == 0] = seats[grid == 0] |
| 63 | + occ = (grid == 2).sum() |
| 64 | + if occ == prev_occ: |
| 65 | + break |
| 66 | + |
| 67 | + print(occ) |
0 commit comments