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))) |