AOC 9 pt 2

day9pt2.py

from functools import reduce
filename = '9.txt'
hmap = []

RIGHT_END = None
with open(filename, 'r') as f:
    for line in f:
        numline = [int(n) for n in list(line.strip())]
        numline = [[n, True] for n in numline]
        RIGHT_END = len(numline) - 1
        hmap.append(numline)

def check_positions(row, col, bottom_row, right_col):
    potential_cells = [
        [row - 1, col],
        [row + 1, col],
        [row, col - 1],
        [row, col + 1],
    ]
    return [c for c in potential_cells if (c[0] >= 0 and c[0] <= bottom_row and c[1] >= 0 and c[1] <= right_col)]

def non_nine_positions(hmap, row, col, bottom_row, right_col):
    potentials = check_positions(row, col, bottom_row, right_col)
    ret = []
    for p in potentials:
        val = hmap[p[0]][p[1]]
        if val[0] != 9:
            ret.append(p)
    return ret

bottom_line = len(hmap) - 1
for line_num in range(len(hmap)):
    right_end = len(hmap[line_num]) - 1
    for col_num in range(len(hmap[line_num])):
        val, is_low = hmap[line_num][col_num]
        if is_low is True:
            check_pos = check_positions(line_num, col_num, bottom_line, right_end)
            is_lowest = True
            for pos in check_pos:
                check_val, _ = hmap[pos[0]][pos[1]]

                if check_val > val:
                    hmap[pos[0]][pos[1]][1] = False
                else:
                    is_lowest = False
            if is_lowest is False:
                hmap[line_num][col_num][1] = False


basins = []
for i, line in enumerate(hmap):
    for j, t in enumerate(line):
        if t[1]:
            basins.append((i, j))

def spread(y, x, hmap, seen=None):
    seen = seen or set()
    if (y, x) not in seen:
        seen.add((y, x))
    for pos in non_nine_positions(hmap, y, x, bottom_line, RIGHT_END):
        if tuple(pos) not in seen:
            seen = seen.union(spread(pos[0], pos[1], hmap, seen))
    return seen

basin_sizes = [0, 0, 0]
for b in basins:
    area = len(spread(b[0], b[1], hmap))
    if area > basin_sizes[0]:
        basin_sizes[0] = area
        basin_sizes.sort()
print(reduce(lambda a, b: a * b, basin_sizes))