day 4: more improvements

1 files changed, 9 insertions, 1 deletions

diff --git a/4.py b/4.py
index f2c3063..0c7cf4f 100644
--- a/4.py
+++ b/4.py
@@ -1,3 +1,4 @@
+from collections import defaultdict
 from utils import open_day
 
 class Board:
@@ -35,13 +36,20 @@ class Board:
 
 def solve(nums: list[int], boards: list[Board]) -> tuple[int, int]:
     indices: set[int] = { i for i in range(len(boards)) }
+    nums_indices: defaultdict[set[int]] = defaultdict(set)
+    i: int
+    board: Board
+    for i, board in enumerate(boards):
+        for n in board.nums:
+            nums_indices[n].add(i)
     wins: list[int] = []
     num: int
     for num in nums:
         i: int
         board: Board
         winners: set[int] = set()
-        for i in indices:
+        num_indices: set[int] = nums_indices.get(num)
+        for i in indices & num_indices:
             boards[i].call(num)
             if boards[i].has_bingo:
                 wins.append(num * boards[i].unmarked_sum)