69 lines
1.4 KiB
Python
69 lines
1.4 KiB
Python
hmap = {}
|
|
with open("input.txt") as f:
|
|
lines = f.read().split('\n')
|
|
height = len(lines)
|
|
width = len(lines[0])
|
|
for i, line in enumerate(lines):
|
|
for j, n in enumerate(line):
|
|
hmap[(i,j)] = int(n)
|
|
|
|
def adjacent(i,j):
|
|
adjacent = [
|
|
(i-1,j),
|
|
(i+1,j),
|
|
(i,j-1),
|
|
(i,j+1)
|
|
]
|
|
return [ (i,j) for (i,j) in adjacent if (-1 < i < height) and (-1 < j < width) ]
|
|
|
|
# part 1
|
|
lowpoints = []
|
|
for i in range(height):
|
|
for j in range(width):
|
|
current = hmap[(i,j)]
|
|
neighbors = [hmap[(x,y)] for (x,y) in adjacent(i,j)]
|
|
if current < min(neighbors):
|
|
lowpoints.append(current)
|
|
risk = 0
|
|
for h in lowpoints:
|
|
risk += (h + 1)
|
|
print(risk)
|
|
|
|
# part 2
|
|
basins = []
|
|
allpoints = set(
|
|
[ (i,j)
|
|
for i in range(height)
|
|
for j in range(width)
|
|
]
|
|
)
|
|
def which_basin(i,j,hmap):
|
|
"""
|
|
Identifies the basin that (i,j) is part of.
|
|
"""
|
|
basin = set()
|
|
visited = set()
|
|
q = [(i,j)]
|
|
|
|
while q:
|
|
(x, y) = q.pop()
|
|
if (x,y) not in visited:
|
|
visited.add((x,y))
|
|
if hmap[(x,y)] != 9:
|
|
basin.add((x,y))
|
|
neighbors = []
|
|
q += [(m,n) for (m,n) in adjacent(x,y) if (m,n) not in visited]
|
|
return basin
|
|
|
|
identified = set()
|
|
for i in range(height):
|
|
for j in range(width):
|
|
if (i,j) not in identified:
|
|
basin = which_basin(i,j,hmap)
|
|
for point in basin:
|
|
identified.add(point)
|
|
basins.append(basin)
|
|
three_biggest_basins = sorted([len(basin) for basin in basins], reverse=True)[0:3]
|
|
from math import prod
|
|
part2 = prod(three_biggest_basins)
|
|
print(part2) |