day 4: switch to a dict of numbers to positions

1 files changed, 12 insertions, 18 deletions

diff --git a/4.py b/4.py
index 1373ca3..c26629b 100644
--- a/4.py
+++ b/4.py
@@ -1,40 +1,34 @@
-from itertools import chain
 from utils import open_day
 
 class Board:
-    cells: list[int]
     width: int
     height: int
     col_hits: list[int]
     row_hits: list[int]
     has_bingo: bool
     unmarked_sum: int
-    nums: set[int]
+    nums: dict[int, tuple[int, int]]
     def __init__(self, cells: list[list[int]]) -> None:
         self.width = len(cells[0])
         self.height = len(cells)
         self.col_hits = [0] * self.width
         self.row_hits = [0] * self.height
-        self.unmarked_sum = 0
-        self.cells = []
-        self.nums = set()
-        for cell in chain.from_iterable(cells):
-            self.unmarked_sum += cell
-            self.cells.append(cell)
-            self.nums.add(cell)
+        self.nums = dict()
+        for y, row in enumerate(cells):
+            for x, cell in enumerate(row):
+                self.nums[cell] = (x, y)
+        self.unmarked_sum = sum(self.nums.keys())
         self.has_bingo = False
     def call(self, num: int) -> None:
-        if num not in self.nums: return
-        pos: int = self.cells.index(num)
-        x: int = pos % self.width
-        y: int = pos // self.width
-        self.col_hits[x] += 1
-        self.row_hits[y] += 1
+        pos: tuple[int, int] = self.nums.get(num)
+        if pos is None: return
+        self.col_hits[pos[0]] += 1
+        self.row_hits[pos[1]] += 1
         self.unmarked_sum -= num
         self.has_bingo = (
             self.has_bingo or
-            self.col_hits[x] == self.height or
-            self.row_hits[y] == self.width
+            self.col_hits[pos[0]] == self.height or
+            self.row_hits[pos[1]] == self.width
         )
     @staticmethod
     def from_string(s: str) -> 'Board':