82 lines
2.3 KiB
Python
82 lines
2.3 KiB
Python
import math
|
|
|
|
VERBOSE = True
|
|
def verbose(s = "", *args, **kwargs):
|
|
if VERBOSE:
|
|
print(s, *args, **kwargs)
|
|
|
|
def get_input(sample = False, part = 1):
|
|
with open(f'sample_p{part}' if sample else 'input', 'r') as f:
|
|
return [[int(c) for c in line.strip()] for line in f.readlines()]
|
|
|
|
|
|
def result(inp, part = 1):
|
|
res = 0;
|
|
width = len(inp[0])
|
|
height = len(inp)
|
|
lowest_h = [0 for _ in range(height)]
|
|
lowest_v = [0 for _ in range(height)]
|
|
lowest = [0 for _ in range(height)]
|
|
|
|
# get lowest in horizontal
|
|
for y, row in enumerate(inp):
|
|
for x, cell in enumerate(row):
|
|
if (x == 0 and cell < row[x+1]) \
|
|
or (x == width - 1 and cell < row[x-1]) \
|
|
or (cell < row[x-1] and cell < row[x+1]):
|
|
lowest_h[y] |= 1 << (width - x)
|
|
|
|
# get lowest in vertical
|
|
for x in range(width):
|
|
for y in range(height):
|
|
row = inp[y]
|
|
cell = row[x]
|
|
if (y == 0 and cell < inp[y+1][x]) \
|
|
or (y == height - 1 and cell < inp[y-1][x]) \
|
|
or (cell < inp[y-1][x] and cell < inp[y+1][x]):
|
|
lowest_v[y] |= 1 << (width - x)
|
|
|
|
# combine both lowests and get the overall lowests
|
|
for y in range(height):
|
|
lowest[y] = lowest_h[y] & lowest_v[y]
|
|
|
|
|
|
lowest_points = []
|
|
for y, row in enumerate(inp):
|
|
for x, cell in enumerate(row):
|
|
if lowest[y] >> (width - x) & 1 == 1:
|
|
if part == 1:
|
|
res += cell + 1
|
|
else:
|
|
lowest_points.append((x, y))
|
|
|
|
if part == 1:
|
|
return res
|
|
|
|
res = 1
|
|
verbose(lowest_points)
|
|
basins = []
|
|
for point in lowest_points:
|
|
queue = [point]
|
|
idx = 0
|
|
while idx < len(queue):
|
|
x, y = queue[idx]
|
|
for d in [(x+1, y), (x-1, y), (x, y+1), (x, y-1)]:
|
|
if d[0] >= 0 and d[0] < width and d[1] >= 0 and d[1] < height \
|
|
and d not in queue and inp[d[1]][d[0]] != 9:
|
|
queue.append(d)
|
|
|
|
idx += 1;
|
|
verbose(f"From { point = }, the region is {len(queue)}")
|
|
basins.append(len(queue))
|
|
|
|
basins.sort()
|
|
res = math.prod(basins[-3:])
|
|
|
|
return res
|
|
|
|
|
|
if __name__ == "__main__":
|
|
VERBOSE = False
|
|
print(result(get_input(), part = 2))
|