🏜

y2022/p14.py

+def build_grid(text):
+    grid = {}
+    for ln in text.splitlines():
+        ns = map(int, ''.join(n if n.isdigit() else ' ' for n in ln).split())
+        nodes = [complex(x, y) for x, y in zip(*[ns] * 2)]
+        for aa, bb in zip(nodes, nodes[1:]):
+            step = (bb - aa) / abs(bb - aa)
+            while aa != bb:
+                grid[aa] = '#'
+                aa += step
+            grid[aa] = '#'
+    return grid
+
+
+def drip(grid, start, floor, p1=False, dof=[1j, 1j - 1, 1j + 1]):
+    pos = start
+    while True:
+        for new in (pos + step for step in dof):
+            if new not in grid and new.imag != floor:
+                pos = new
+                break
+        else:
+            grid[pos] = 'o'
+            yield True
+            pos = start
+
+        if p1:
+            if pos.imag == floor - 1:
+                break
+        else:
+            if grid.get(start) == 'o':
+                break
+
+
+def main():
+    text = open(0).read()
+    grid = build_grid(text)
+    start = complex(500, 0)
+    floor = max(p.imag for p in grid) + 2
+
+    ans1 = sum(drip(grid.copy(), start, floor, True))
+    print(ans1)
+
+    ans2 = sum(drip(grid.copy(), start, floor))
+    print(ans2)
+
+
+main()