Spaces:
Running
Running
| from collections import defaultdict | |
| from typing import List, Set, Tuple, Dict | |
| def read_grid(filename: str) -> List[List[int]]: | |
| with open(filename) as f: | |
| return [[int(c) for c in line.strip()] for line in f] | |
| def get_neighbors(grid: List[List[int]], x: int, y: int) -> List[Tuple[int, int]]: | |
| directions = [(0, 1), (1, 0), (0, -1), (-1, 0)] | |
| height, width = len(grid), len(grid[0]) | |
| neighbors = [] | |
| for dx, dy in directions: | |
| nx, ny = x + dx, y + dy | |
| if 0 <= nx < height and 0 <= ny < width: | |
| neighbors.append((nx, ny)) | |
| return neighbors | |
| def find_trailheads(grid: List[List[int]]) -> List[Tuple[int, int]]: | |
| return [(i, j) for i in range(len(grid)) for j in range(len(grid[0])) if grid[i][j] == 0] | |
| def count_reachable_nines(grid: List[List[int]], start: Tuple[int, int]) -> int: | |
| def dfs(x: int, y: int, visited: Set[Tuple[int, int]]) -> Set[Tuple[int, int]]: | |
| nines = set() | |
| visited.add((x, y)) | |
| current_height = grid[x][y] | |
| if current_height == 9: | |
| nines.add((x, y)) | |
| for nx, ny in get_neighbors(grid, x, y): | |
| if (nx, ny) not in visited and grid[nx][ny] == current_height + 1: | |
| nines.update(dfs(nx, ny, visited)) | |
| return nines | |
| return len(dfs(start[0], start[1], set())) | |
| def count_distinct_paths(grid: List[List[int]], start: Tuple[int, int]) -> int: | |
| memo: Dict[Tuple[int, int], int] = {} | |
| def dfs(x: int, y: int) -> int: | |
| if (x, y) in memo: | |
| return memo[(x, y)] | |
| if grid[x][y] == 9: | |
| return 1 | |
| paths = 0 | |
| current_height = grid[x][y] | |
| for nx, ny in get_neighbors(grid, x, y): | |
| if grid[nx][ny] == current_height + 1: | |
| paths += dfs(nx, ny) | |
| memo[(x, y)] = paths | |
| return paths | |
| return dfs(start[0], start[1]) | |
| def solve_part1(grid: List[List[int]]) -> int: | |
| trailheads = find_trailheads(grid) | |
| return sum(count_reachable_nines(grid, start) for start in trailheads) | |
| def solve_part2(grid: List[List[int]]) -> int: | |
| trailheads = find_trailheads(grid) | |
| return sum(count_distinct_paths(grid, start) for start in trailheads) | |
| grid = read_grid("./input.txt") | |
| print(str(solve_part1(grid))) | |
| print(str(solve_part2(grid))) |