Day 11 (part 2 more beautiful / method of part 1 is faster)

day11/part2.py

@@ -1,50 +1,42 @@
 import numpy as np
 from day11.common import char_to_int
 
-lines = np.array([[char_to_int(x) for x in line.strip()] for line in open("input.txt")], dtype=np.byte)
-shape = lines.shape
-
+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)]
+occupied = -1
+
+
+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
+        elif seats[idx] == 0:
+            idx = tuple(map(sum, zip(idx, ray)))
+            continue
+        else:
+            return 1
+    else:
+        return 0
+
 
-last = -1
 while True:
-    curr = np.count_nonzero(lines == 1)
-    if curr == last:
-        print(last)
+    new_occupied = np.count_nonzero(seats == 1)
+    if occupied == new_occupied:
         break
-    last = curr
-
-    new = np.zeros_like(lines, dtype=np.byte)
-    for idx, val in np.ndenumerate(lines):
+    occupied = new_occupied
+    result = np.zeros_like(seats)
+    for idx, val in np.ndenumerate(seats):
         if val == 0:
             continue
-        if val == -1:
-            found = False
-            for ray in RAYS:
-                ray = np.array(ray)
-                pos = np.array(idx) + ray
-                while 0 <= pos[0] < shape[0] and 0 <= pos[1] < shape[1]:
-                    if lines[tuple(pos)] == 1:
-                        found = True
-                        break
-                    elif lines[tuple(pos)] == -1:
-                        break
-                    pos = pos + ray
-                if found:
-                    break
-            new[idx] = 1 if not found else -1
-            continue
-        if val == 1:
-            found = 0
-            for ray in RAYS:
-                ray = np.array(ray)
-                pos = np.array(idx) - ray
-                while 0 <= pos[0] < shape[0] and 0 <= pos[1] < shape[1]:
-                    if lines[tuple(pos)] == 1:
-                        found += 1
-                        break
-                    elif lines[tuple(pos)] == -1:
-                        break
-                    pos = pos - ray
-            new[idx] = -1 if found >= 5 else 1
-    lines = new
+        neighbors = sum(find_neighbor(idx, ray) for ray in RAYS)
+        if val == -1 and neighbors == 0:
+            result[idx] = 1
+        elif val == 1 and neighbors >= 5:
+            result[idx] = -1
+        else:
+            result[idx] = seats[idx]
+    seats = result
+    print(occupied)
+
+print(occupied)