advent_of_code_2021/day09/part02.py

import argparse
import functools
from typing import NamedTuple, List, Dict, Set, Tuple

parser = argparse.ArgumentParser()
parser.add_argument("ifile", type=argparse.FileType('r'))
args = parser.parse_args()

heightmap = [[int(x) for x in line.strip()] for line in args.ifile.readlines()]


def printMap(map: List[List]):
    for line in map:
        print("".join([str(int(el)) for el in line]))

def printColorMap(map: List[List]):
    for line in map:
        print("".join([el for el in line]))

def descendingMap(seq: List[int]) -> List[bool]:
    return [True] + list(map(lambda pair: pair[0] < pair[1], zip(seq[1:], seq)))


right = map(descendingMap, heightmap)
left = map(lambda line: reversed(descendingMap(list(reversed(line)))), heightmap)
horiz = [[all(pair) for pair in line] for line in map(lambda pair: zip(pair[0], pair[1]), zip(right, left))]

# transpose array to check cols
flipped = [[heightmap[j][i] for j in range(len(heightmap))] for i in range(len(heightmap[0]))]
up = map(descendingMap, flipped)
down = map(lambda line: reversed(descendingMap(list(reversed(line)))), flipped)
pre_vert = [[all(pair) for pair in line] for line in map(lambda pair: zip(pair[0], pair[1]), zip(up, down))]

# transpose back
vert = [[pre_vert[j][i] for j in range(len(pre_vert))] for i in range(len(pre_vert[0]))]
composite = [[all(pair) for pair in line] for line in map(lambda pair: zip(pair[0], pair[1]), zip(horiz, vert))]
paired = [[pair[0] + 1 if pair[1] else 0 for pair in line] for line in map(lambda pair: zip(pair[0], pair[1]), zip(heightmap, composite))]


def neighborsForPos(x: int, y: int):
    return list(filter(lambda pair: 0 <= pair[0] and pair[0] < len(heightmap[0]) and 0 <= pair[1] and pair[1] < len(heightmap),
                       [(x - 1, y), (x, y + 1), (x + 1, y), (x, y - 1)]))


# printMap(heightmap)
# print()
# printMap(composite)
# print()

def search(start_x: int, start_y: int, been: set[Tuple[int, int]]) -> int:
    been.add((start_x, start_y))
    our_val = heightmap[start_y][start_x]

    for neighbor in filter(lambda neighbor: neighbor not in been, neighborsForPos(start_x, start_y)):
        neighbor_val = heightmap[neighbor[1]][neighbor[0]]
        if (neighbor_val != 9) and (neighbor_val > our_val):
            search(neighbor[0], neighbor[1], been)

basins: List[Set[Tuple[int, int]]] = []
for y in range(0, len(heightmap)):
    for x in range(0, len(heightmap[0])):
        if composite[y][x]:
            been: Set[Tuple[int, int]] = set()
            search(x, y, been)
            basins.append(been)

# colors = "ABCDEFGHIJK"

# colormap = [["." for _ in range(0, len(heightmap[0]))] for _ in heightmap]

# for index, basin in enumerate(basins):
#     for point in basin:
#         colormap[point[1]][point[0]] = colors[index]
# print(basins)
# printColorMap(colormap)

basin_sizes = list(map(len, basins))
basin_sizes.sort()
print(functools.reduce(lambda acc, val: acc * val, basin_sizes[::-1][:3]))