9_1.py

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