1 files changed, 66 insertions, 0 deletions
diff --git a/2021/09/puzzle-2.py b/2021/09/puzzle-2.py
new file mode 100755
index 0000000..d435f7c
--- /dev/null
+++ b/2021/09/puzzle-2.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+
+"""
+Most of the code taken from this guys floodfill, I just adapted it for this problem
+https://playandlearntocode.com/article/flood-fill-algorithm-in-python
+"""
+
+from functools import reduce
+from itertools import product
+from operator import mul
+
+matrix = list[list[int]]
+
+
+def is_basin(grid: matrix, row: int, col: int) -> bool:
+	return (
+		not ((row < 0 or row > len(grid) - 1) or (col < 0 or col > len(grid[0]) - 1))
+		and grid[row][col] != 9
+	)
+
+
+def floodfill(grid: matrix, row: int, col: int) -> int:
+	if row < 0 or row > len(grid) - 1 or col < 0 or col > len(grid[0]) - 1 or grid[row][col] == 9:
+		return 0
+
+	q: list[tuple[int, int]] = []
+	grid[row][col] = 9
+	q.append((row, col))
+	size = 1
+
+	while len(q):
+		row, col = q.pop(0)
+
+		for x, y in ((-1, 0), (1, 0), (0, -1), (0, 1)):
+			c_row = row + x
+			c_col = col + y
+
+			if is_basin(grid, c_row, c_col):
+				grid[c_row][c_col] = 9
+				q.append((c_row, c_col))
+				size += 1
+
+	return size
+
+
+def solve(grid: matrix) -> int:
+	return reduce(
+		mul,
+		sorted(
+			map(
+				lambda t: floodfill(grid, t[0], t[1]),
+				filter(
+					lambda t: grid[t[0]][t[1]] != 9, product(range(len(grid)), range(len(grid[0]))),
+				),
+			)
+		)[-3:],
+	)
+
+
+def main() -> None:
+	with open("input", "r", encoding="utf-8") as f:
+		print(solve(list(map(lambda l: [int(n) for n in l.strip()], f.readlines()))))
+
+
+if __name__ == "__main__":
+	main()