81 lines
2 KiB
Python
81 lines
2 KiB
Python
#!/usr/bin/env python3
|
|
|
|
trail_map = []
|
|
|
|
with open("day10") as f:
|
|
for l in f:
|
|
trail_map.append(list(map(int, l.strip())))
|
|
|
|
def trail_heads(trail_map):
|
|
heads = []
|
|
for y,_ in enumerate(trail_map):
|
|
for x,_ in enumerate(trail_map[y]):
|
|
if trail_map[y][x] == 0:
|
|
heads.append((y,x))
|
|
return heads
|
|
|
|
from copy import copy
|
|
|
|
def check_path_dfs(head, trail_map, part2):
|
|
visited = set()
|
|
stack = [(head, [head])]
|
|
|
|
def is_height(y,x,h):
|
|
if x<0 or x>=len(trail_map[0]):
|
|
return False
|
|
if y<0 or y>=len(trail_map[0]):
|
|
return False
|
|
else:
|
|
return trail_map[y][x] == h
|
|
|
|
trails = []
|
|
visited_trails = []
|
|
while stack:
|
|
el_tr = stack.pop()
|
|
el = el_tr[0]
|
|
tr = el_tr[1]
|
|
|
|
if part2 and tr in visited_trails:
|
|
continue
|
|
|
|
if part2: visited_trails.append(tr)
|
|
if not part2: visited.add(el)
|
|
|
|
el_height = trail_map[el[0]][el[1]]
|
|
|
|
if el_height == 9:
|
|
trails.append(tr)
|
|
continue
|
|
|
|
up = (el[0]-1 , el[1])
|
|
down = (el[0]+1 , el[1])
|
|
left = (el[0] , el[1]-1)
|
|
right = (el[0] , el[1]+1)
|
|
|
|
up_ok = up not in visited and is_height(*up, el_height+1)
|
|
down_ok = down not in visited and is_height(*down, el_height+1)
|
|
left_ok = left not in visited and is_height(*left, el_height+1)
|
|
right_ok = right not in visited and is_height(*right, el_height+1)
|
|
|
|
if up_ok: stack.append( (up, copy(tr)+[up]) )
|
|
if down_ok: stack.append( (down, copy(tr)+[down]) )
|
|
if left_ok: stack.append( (left, copy(tr)+[left]) )
|
|
if right_ok: stack.append( (right, copy(tr)+[right]) )
|
|
|
|
return trails
|
|
|
|
|
|
heads = trail_heads(trail_map)
|
|
|
|
res = 0
|
|
for head in heads:
|
|
paths = check_path_dfs(head, trail_map, part2=False)
|
|
res += len(paths)
|
|
print(res)
|
|
|
|
res = 0
|
|
for head in heads:
|
|
paths = check_path_dfs(head, trail_map, part2=True)
|
|
res += len(paths)
|
|
print(res)
|
|
|