2021 day9/python: add cleaner solution

2021/day9/clean.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+
+from collections import Counter
+from dataclasses import dataclass
+import math
+import sys
+from typing import Optional, Tuple
+
+@dataclass
+class Point:
+    height: int
+    low_point: Optional[Tuple[int, int]]
+
+class Grid:
+    def __init__(self, data):
+        self.data = data
+        self.rows = len(data)
+        self.cols = len(data[0])
+
+    def get(self, x, y):
+        if 0 <= x < self.rows and 0 <= y < self.cols:
+            return self.data[x][y]
+        else:
+            return None
+
+    def find_low_point(self, x, y):
+        current = self.get(x, y)
+        if current is None:
+            raise ValueError
+
+        if current.low_point is not None:
+            return current.low_point
+        if current.height == 9:
+            return None
+
+        for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)):
+            target = self.get(x+dx, y+dy)
+            if target is not None:
+                if target.height < current.height:
+                    current.low_point = self.find_low_point(x+dx, y+dy)
+                    return current.low_point
+        return (x, y)
+
+
+if __name__ == '__main__':
+    grid = Grid([
+        [
+            Point(height=int(n), low_point=None)
+            for n in i.strip()
+        ]
+        for i in sys.stdin.readlines()
+    ])
+
+    counts = Counter()
+    for x in range(grid.rows):
+        for y in range(grid.cols):
+            low_point = grid.find_low_point(x, y)
+            if low_point is not None:
+                counts.update([low_point])
+
+    print(sum(grid.get(x, y).height + 1 for x, y in counts.keys()))
+    print(math.prod(c for _, c in counts.most_common(3)))