AOC day 10 P1&p2

aoc_d1_p1.py

visited_trail_heads = {}
starters = []
grid = []

for line in open('input10.txt'):
    grid.append(line)
n = len(grid)

for i in range(n):
    for j in range(n):
        if grid[i][j] == '0':
            starters.append((i, j))
        elif grid[i][j] == '9':
            if (i, j) not in visited_trail_heads:
                visited_trail_heads[(i, j)] = set()

dirs = [(1, 0), (-1, 0), (0, 1), (0, -1)]

dir_id = 0


def dfs(x, y, s, direction):
    global dir_id
    if grid[x][y] == '9':
        # for P1, just remove the direction component.
        visited_trail_heads[(x, y)].add((s, direction))
        return

    for d in dirs:
        new_x = x + d[0]
        new_y = y + d[1]
        if 0 <= new_x < n and 0 <= new_y < n and grid[new_x][new_y] != '.' and int(grid[new_x][new_y]) == int(
                grid[x][y]) + 1:
            dir_id += 1
            dfs(new_x, new_y, s, dir_id)


for s in starters:
    dfs(s[0], s[1], (s[0], s[1]), -1)

ans = 0
print(visited_trail_heads)
for s in visited_trail_heads:
    ans += len(visited_trail_heads[s])
print(ans)