advent_of_code/9/9_2.py

#!/usr/bin/env python
# 2021 - Advent Of Code - 9 part 2

def parse_file(file):
    hmap = []
    with open(file) as f:
        for line in f.readlines():
            linemap = []
            for n in line.strip():
                linemap.append(int(n))
            hmap.append(linemap)

    return hmap


def find_lowpoints(hmap):
    lpoints = []

    # first line
    x = 0
    y = 0
    if hmap[y][x] < hmap[y][x + 1] and hmap[y][x] < hmap[y + 1][x]:
        lpoints.append([x, y])

    for x in range(1, len(hmap[y]) - 1):
        v = hmap[y][x]
        if v < hmap[y][x - 1] and v < hmap[y][x + 1] and v < hmap[y + 1][x]:
            lpoints.append([x, y])

    x = len(hmap[y]) - 1
    if hmap[y][x] < hmap[y][x - 1] and hmap[y][x] < hmap[y + 1][x]:
        lpoints.append([x, y])

    # middle lines
    for y in range(1, len(hmap) - 1):
        x = 0
        if hmap[y][x] < hmap[y][x + 1] and hmap[y][x] < hmap[y - 1][x] and hmap[y][x] < hmap[y + 1][x]:
            lpoints.append([x, y])

        for x in range(1, len(hmap[y]) - 1):
            v = hmap[y][x]
            if v < hmap[y][x - 1] and v < hmap[y][x + 1] and v < hmap[y - 1][x] and v < hmap[y + 1][x]:
                lpoints.append([x, y])

        x = len(hmap[y]) - 1
        if hmap[y][x] < hmap[y][x - 1] and hmap[y][x] < hmap[y - 1][x] and hmap[y][x] < hmap[y + 1][x]:
            lpoints.append([x, y])

    # last line
    y = len(hmap) - 1
    x = 0
    if hmap[y][x] < hmap[y][x + 1] and hmap[y][x] < hmap[y - 1][x]:
        lpoints.append([x, y])

    for x in range(1, len(hmap[y]) - 1):
        v = hmap[y][x]
        if v < hmap[y][x - 1] and v < hmap[y][x + 1] and v < hmap[y - 1][x]:
            lpoints.append([x, y])

    x = len(hmap[y]) - 1
    if hmap[y][x] < hmap[y][x - 1] and hmap[y][x] < hmap[y - 1][x]:
        lpoints.append([x, y])

    return lpoints


def find_basin(hmap, bpoints, point):
    bpoints.append(point)
    # print(f'bpoints: {bpoints}')
    x = point[0]
    y = point[1]

    # halt if edge, already in list or value is 9
    # up
    if y > 0:
        if hmap[y-1][x] != 9:
            if [x, y-1] not in bpoints:
                find_basin(hmap, bpoints, [x, y-1])
    # down
    if y < len(hmap) - 1:
        if hmap[y+1][x] != 9:
            if [x, y+1] not in bpoints:
                find_basin(hmap, bpoints, [x, y+1])

    # left
    if x > 0:
        if hmap[y][x-1] != 9:
            if [x-1, y] not in bpoints:
                find_basin(hmap, bpoints, [x-1, y])

    # right
    if x < len(hmap[y]) - 1:
        if hmap[y][x+1] != 9:
            if [x+1, y] not in bpoints:
                find_basin(hmap, bpoints, [x+1, y])

    return bpoints


# heightmap = parse_file('input_example.txt')
heightmap = parse_file('input.txt')
# print_hmap(heightmap)
low_points = find_lowpoints(heightmap)

risk_levels = 0
basins = []
for p in low_points:
    basins_points = []
    find_basin(heightmap, basins_points, p)
    basins.append(basins_points)


sizes = [0] * len(basins)
i = 0
for b in basins:
    size = 0
    for p in b:
        size += 1
    sizes[i] = size
    i += 1

# 3 biggest basins
biggest_basins = sorted(zip(sizes, basins))[-3:]

product = 1
for s, b in biggest_basins:
    print(f'basin size: {s}')
    product *= s

print(f'product: {product}')