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