Advent of Code 2021 - day 9

aoc21_9.py

grid = []
with open('input.txt', 'r') as f:
    for line in f.readlines():
        grid.append([int(x) for x in line.strip()])
height = len(grid)
width = len(grid[0])

def get_nbrs(i, j):
    for di, dj in ((-1, 0), (1, 0), (0, -1), (0, 1)):
        nbri, nbrj = i+di, j+dj
        if 0 <= nbri < height and 0 <= nbrj < width:
            yield nbri, nbrj

def get_basin_size(i, j):
    # get size of basin for given element in basin
    start = (i, j)
    q = [start]
    visited = {start}
    while q:
        current = q.pop()
        for nbri, nbrj in get_nbrs(current[0], current[1]):
            if grid[nbri][nbrj] == 9:
                continue
            nbr = (nbri, nbrj)
            if nbr in visited:
                continue
            q.append(nbr)
            visited.add(nbr)
    return len(visited)

result1 = 0
basin_sizes = []
for i in range(height):
    for j in range(width):
        h = grid[i][j]
        if all(grid[nbri][nbrj] > h for (nbri, nbrj) in get_nbrs(i, j)):
            result1 += (h+1)
            basin_sizes.append(get_basin_size(i, j))
# part 1
print(result1)
# part 2
basin_sizes.sort(reverse=True)
print(basin_sizes[0]*basin_sizes[1]*basin_sizes[2])