utils: Point2D and related functions

1 files changed, 21 insertions, 23 deletions

diff --git a/9.py b/9.py
index d1a63b9..0db75b1 100644
--- a/9.py
+++ b/9.py
@@ -1,45 +1,43 @@
 from functools import reduce
-from utils import open_day
+from utils import open_day, Point2D, adjacent_bounded
 from collections import deque
 
 heatmap = [[int(c) for c in l.rstrip()] for l in open_day(9)]
+bounds = Point2D(len(heatmap[0]), len(heatmap))
 adj = [(1, 0), (0, 1), (-1, 0), (0, -1)]
 
-def is_low(x, y):
-    risk = heatmap[y][x]
-    for dx, dy in adj:
-        nx = x + dx
-        ny = y + dy
-        if nx >= 0 and nx < len(heatmap[0]) and ny >= 0 and ny < len(heatmap):
-            if heatmap[y+dy][x+dx] <= risk:
-                return False
+def is_low(p):
+    risk = heatmap[p.y][p.x]
+    for n in adjacent_bounded(p, bounds, False):
+        if heatmap[n.y][n.x] <= risk:
+            return False
     return True
 
-def basin_size(x, y):
+def basin_size(p):
     seen = set()
-    queue = deque(((x, y),))
+    queue = deque((p,))
     size = 0
     while queue:
-        x, y = queue.popleft()
-        if (x, y) in seen:
+        p = queue.popleft()
+        if p in seen:
             continue
-        if x < 0 or x >= len(heatmap[0]) or y < 0 or y >= len(heatmap):
-            continue
-        seen.add((x, y))
-        if heatmap[y][x] == 9:
+        seen.add(p)
+        if heatmap[p.y][p.x] == 9:
             continue
         size += 1
-        for dx, dy in adj:
-            queue.append((x + dx, y + dy))
+        for n in adjacent_bounded(p, bounds, False):
+            # print(p, n)
+            queue.append(n)
     return size
 
 totrisk = 0
 lows = []
-for y in range(len(heatmap)):
-    for x in range(len(heatmap[y])):
-        if is_low(x, y):
+for y in range(bounds.y):
+    for x in range(bounds.x):
+        p = Point2D(x, y)
+        if is_low(p):
             totrisk += 1 + heatmap[y][x]
-            lows.append(basin_size(x, y))
+            lows.append(basin_size(p))
 print(totrisk)
 lows.sort()
 print(reduce(lambda x, y: x * y, lows[-3:], 1))