day09_2.py

from functools import reduce


def calc_basin_size(y, x, matrix, fill_matrix, size):
    fill_matrix[y][x] = 1
    up_size = 0
    dn_size = 0
    lf_size = 0
    rt_size = 0
    if y > 0 and fill_matrix[y - 1][x] == 0 and matrix[y-1][x] != 9:
        up_size = calc_basin_size(y - 1, x, matrix, fill_matrix, size)
    if y < len(matrix)-1 and fill_matrix[y + 1][x] == 0 and matrix[y+1][x] != 9:
        dn_size = calc_basin_size(y + 1, x, matrix, fill_matrix, size)
    if x > 0 and fill_matrix[y][x - 1] == 0 and matrix[y][x-1] != 9:
        lf_size = calc_basin_size(y, x - 1, matrix, fill_matrix, size)
    if x < len(matrix[y])-1 and fill_matrix[y][x + 1] == 0 and matrix[y][x+1] != 9:
        rt_size = calc_basin_size(y, x + 1, matrix, fill_matrix, size)
    return size + up_size + dn_size + lf_size + rt_size


def solve(lines):
    matrix = []
    fill_matrix = []
    for line in lines:
        matrix.append(list(map(int, [c for c in line])))
        fill_matrix.append([0] * len(line))

    basin_sizes = []
    for y in range(0, len(lines)):
        for x in range(0, len(lines[y])):
            up = 10 if y == 0 else int(lines[y - 1][x])
            down = 10 if y == len(lines) - 1 else int(lines[y + 1][x])
            left = 10 if x == 0 else int(lines[y][x - 1])
            right = 10 if x == len(lines[y]) - 1 else int(lines[y][x + 1])
            point = int(lines[y][x])
            if point < up and point < down and point < left and point < right:
                basin_sizes.append(calc_basin_size(y, x, matrix, fill_matrix, 1))
    return reduce((lambda a, b: a * b), sorted(basin_sizes, reverse=True)[:3])


if __name__ == "__main__":
    with open("input.txt") as file:
        lines = [line.strip() for line in file.readlines()]

    sum = solve(lines)

    print(sum)

    assert sum == 900864