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]))