vigilantes

y2021/p09.py

@@ -0,0 +1,23 @@
+from toolkit import read_image
+
+
+ans1 = 0
+grid = read_image(text)[0]
+for p, v in grid.items():
+    if all(grid[p + s] > v for s in [1, -1, 1j, -1j] if p + s in grid):
+        ans1 += int(v) + 1
+
+
+valid = {k for k, v in grid.items() if v != '9'}
+pending = valid.copy()
+basin_sizes = []
+while pending:
+    edge = {pending.pop()}
+    seen = edge.copy()
+    while edge:
+        edge = {p + s for p in edge for s in [1, -1, 1j, -1j]} & valid - seen
+        seen |= edge
+    pending -= seen
+    basin_sizes.append(len(seen))
+*_, x, y, z = sorted(basin_sizes)
+ans2 = x * y * z