Solution to problem 9 in Python

src/Year_2021/P9.py

@@ -40,6 +40,7 @@
 
 
 from collections import defaultdict
+from math import prod
 
 with open("files/P9.txt") as f:
     heightmap = [line for line in f.read().strip().split()]
@@ -69,5 +70,87 @@ def part_1() -> None:
     print(f"The sum of all low points is {res}")
 
 
+# --- Part Two ---
+
+# Next, you need to find the largest basins so you know what areas are most
+# important to avoid.
+
+# A basin is all locations that eventually flow downward to a single low point.
+# Therefore, every low point has a basin, although some basins are very small.
+# Locations of height 9 do not count as being in any basin, and all other
+# locations will always be part of exactly one basin.
+
+# The size of a basin is the number of locations within the basin, including
+# the low point. The example above has four basins.
+
+# The top-left basin, size 3:
+
+# 2199943210
+# 3987894921
+# 9856789892
+# 8767896789
+# 9899965678
+
+# The top-right basin, size 9:
+
+# 2199943210
+# 3987894921
+# 9856789892
+# 8767896789
+# 9899965678
+
+# The middle basin, size 14:
+
+# 2199943210
+# 3987894921
+# 9856789892
+# 8767896789
+# 9899965678
+
+# The bottom-right basin, size 9:
+
+# 2199943210
+# 3987894921
+# 9856789892
+# 8767896789
+# 9899965678
+
+# Find the three largest basins and multiply their sizes together. In the above
+# example, this is 9 * 14 * 9 = 1134.
+
+# What do you get if you multiply together the sizes of the three largest
+# basins?
+
+
+def bassin(
+    queue: list[tuple[int, int]], is_visited: set[tuple[int, int]]
+) -> int:
+    while queue:
+        x, y = queue.pop(0)
+        for neighbor in ((x - 1, y), (x + 1, y), (x, y + 1), (x, y - 1)):
+            value = heightmap_dic.get(neighbor, 9)
+            if (
+                value == 9
+                or neighbor in is_visited
+                or heightmap_dic[(x, y)] > value
+            ):
+                continue
+            is_visited.add(neighbor)
+            queue.append(neighbor)
+
+    return len(is_visited)
+
+
+def part_2() -> None:
+    res = []
+    for pos, value in heightmap_dic.items():
+        if not is_low(*pos, value):
+            continue
+        res.append(bassin([pos], {pos}))
+
+    print(prod(sorted(res)[-3:]))
+
+
 if __name__ == "__main__":
     part_1()
+    part_2()