day 4: refactor into a function

1 files changed, 32 insertions, 23 deletions

diff --git a/4.py b/4.py
index 7ad6067..3cafb8b 100644
--- a/4.py
+++ b/4.py
@@ -1,32 +1,41 @@
 from utils import open_day
+from dataclasses import dataclass
 
-nums, *boards = open_day(4).read().split('\n\n')
-nums = list(map(int, nums.split(',')))
-boards = [[[(int(n), False) for n in l.split()] for l in b.rstrip().split('\n')] for b in boards]
+@dataclass
+class Cell:
+    num: int
+    marked: bool = False
 
-def win(board):
-    for y in board:
-        if all(x[1] for x in y):
+def bingo(board):
+    for row in board:
+        if all(cell.marked for cell in row):
             return True
     for x in range(len(board[0])):
-        if all(board[y][x][1] for y in range(len(board))):
+        if all(board[y][x].marked for y in range(len(board))):
             return True
     return False
 
-won = set()
-wins = []
-for n in nums:
-    for i, b in enumerate(boards):
-        tot = 0
-        for y in range(len(b)):
-            for x in range(len(b[y])):
-                if b[y][x][0] == n:
-                    b[y][x] = (b[y][x][0], True)
-                if b[y][x][1] == False:
-                    tot += b[y][x][0]
-        if i not in won and win(b):
-            won.add(i)
-            wins.append(n * tot)
+def solve(nums, boards):
+    won = set()
+    wins = []
+    for num in nums:
+        for i, board in enumerate(boards):
+            total = 0
+            for row in board:
+                for cell in row:
+                    if cell.num == num:
+                        cell.marked = True
+                    if not cell.marked:
+                        total += cell.num
+            if i not in won and bingo(board):
+                won.add(i)
+                wins.append(num * total)
+    return wins[0], wins[-1]
+
+nums, *boards = open_day(4).read().rstrip().split('\n\n')
+nums = list(map(int, nums.split(',')))
+boards = [[[Cell(int(n)) for n in l.split()] for l in b.split('\n')] for b in boards]
 
-print(wins[0])
-print(wins[-1])
+part1, part2 = solve(nums, boards)
+print(part1)
+print(part2)