9_2.py

   1 """Find where smoke settles"""
   2 import math
   3
   4 heights = []
   5
   6 with open("9_input.txt") as f:
   7     for line in f:
   8         heights.append([int(char) for char in line.strip()])
   9
  10 n_rows = len(heights)
  11 n_cols = len(heights[0])
  12
  13 mins = []
  14 mins_coords = []
  15 for i in range(n_rows):
  16     for j in range(n_cols):
  17         val = heights[i][j]
  18         to_compare = [
  19             heights[i - 1][j] if i != 0 else 10,
  20             heights[i][j - 1] if j != 0 else 10,
  21             10 if i == n_rows - 1 else heights[i + 1][j],
  22             10 if j == n_cols - 1 else heights[i][j + 1],
  23         ]
  24         if val < min(to_compare):
  25             mins.append(val)
  26             mins_coords.append((i, j))
  27
  28 basins = []
  29 for minimum in mins_coords:
  30     to_walk = set([minimum])
  31     walked = set()
  32
  33     while to_walk:
  34         i, j = to_walk.pop()
  35         orth_coords = [(i - 1, j), (i, j - 1), (i + 1, j), (i, j + 1)]
  36         for x, y in orth_coords:
  37             try:
  38                 if (x >= 0 and y >= 0 and heights[x][y] != 9 and
  39                         (x, y) not in walked):
  40                     to_walk.add((x, y))
  41             except IndexError:
  42                 pass
  43
  44         walked.add((i, j))
  45
  46     basins.append(walked)
  47
  48 basins = sorted(basins, key=lambda x: len(x))
  49 print(math.prod([len(b) for b in basins[-3:]]))