Day 11 (part 2 more beautiful with cache)

day11/part2.py

@@ -2,7 +2,9 @@ import numpy as np
 from day11.common import char_to_int
 
 seats = np.array([[char_to_int(x) for x in line.strip()] for line in open("input.txt")], dtype=np.byte)
-RAYS = [(-1, -1), (0, -1), (1, -1), (-1, 0), (1, 0), (-1, 1), (0, 1), (1, 1)]
+neighbors = np.zeros_like(seats, dtype=np.byte)
+next_neighbors = np.zeros_like(seats, dtype=np.byte)
+RAYS = [(1, 0), (-1, 1), (0, 1), (1, 1)]
 occupied = -1
 
 
@@ -10,14 +12,14 @@ def find_neighbor(idx, ray):
     idx = tuple(map(sum, zip(idx, ray)))
     while 0 <= idx[0] < seats.shape[0] and 0 <= idx[1] < seats.shape[1]:
         if seats[idx] == -1:
-            return 0
+            return 0, idx
         elif seats[idx] == 0:
             idx = tuple(map(sum, zip(idx, ray)))
             continue
         else:
-            return 1
+            return 1, idx
     else:
-        return 0
+        return None
 
 
 while True:
@@ -29,14 +31,21 @@ while True:
     for idx, val in np.ndenumerate(seats):
         if val == 0:
             continue
-        neighbors = sum(find_neighbor(idx, ray) for ray in RAYS)
-        if val == -1 and neighbors == 0:
+        findings = (find_neighbor(idx, ray) for ray in RAYS)
+        findings = [f for f in findings if f is not None]
+        neighbors[idx] += sum(occ for occ, f in findings)
+        if val == -1 and neighbors[idx] == 0:
             result[idx] = 1
-        elif val == 1 and neighbors >= 5:
+        elif val == 1 and neighbors[idx] >= 5:
             result[idx] = -1
         else:
             result[idx] = seats[idx]
+        if result[idx] == 1:
+            for occ, f in findings:
+                next_neighbors[f] += 1
     seats = result
+    neighbors = next_neighbors
+    next_neighbors = np.zeros_like(seats, dtype=np.byte)
     print(occupied)
 
 print(occupied)