advent_9.py

def is_low_point(heights, x, y):
  adjacent = []
  if x > 0:
    adjacent.append(heights[y][x - 1])
  if x < len(heights[y]) - 1:
    adjacent.append(heights[y][x + 1])
  if y > 0:
    adjacent.append(heights[y - 1][x])
  if y < len(heights) - 1:
    adjacent.append(heights[y + 1][x])
  if heights[y][x] < min(adjacent):
    return True
  return False

def find_risk_levels(heights):
  for y, row in enumerate(heights):
    for x, height in enumerate(row):
      if is_low_point(heights, x, y):
        yield height + 1

def find_basin_sizes(heights):
  visited = set()
  for y, row in enumerate(heights):
    for x, height in enumerate(row):
      basin_size = visit_basin(heights, x, y, visited)
      if basin_size != 0:
        yield basin_size

def visit_basin(heights, x, y, visited):
  def key(x, y):
    return "{0}-{1}".format(x, y)
  if y < 0 or y >= len(heights) or x < 0 or x >= len(heights[y]):
    return 0 # Out of bounds
  current_key = key(x, y)
  if current_key in visited:
    return 0 # Already covered
  if heights[y][x] == 9:
    return 0
  visited.add(current_key)
  return (1 + visit_basin(heights, x + 1, y, visited) +
      visit_basin(heights, x - 1, y, visited) +
      visit_basin(heights, x, y + 1, visited) +
      visit_basin(heights, x, y - 1, visited))

if __name__ == '__main__':
  with open('input.txt') as data:
    heights = []
    for line in data:
      heights.append([int(char) for char in line.strip()])
    print(sum(find_risk_levels(heights)))
    largest_3_basins = sorted(find_basin_sizes(heights))[-3:]
    product = 1
    for item in largest_3_basins:
      product *= item
    print(product)