Spaces:
Running
Running
File size: 2,375 Bytes
a4da721 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
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))) |