Spaces:
Running
Running
| import heapq | |
| def solve(): | |
| file = "input.txt" | |
| grid = [] | |
| with open(file, 'r') as f: | |
| for line in f: | |
| grid.append(line.strip()) | |
| rows = len(grid) | |
| cols = len(grid[0]) | |
| start = None | |
| end = None | |
| for r in range(rows): | |
| for c in range(cols): | |
| if grid[r][c] == 'S': | |
| start = (r, c) | |
| elif grid[r][c] == 'E': | |
| end = (r, c) | |
| directions = [(0, 1), (1, 0), (0, -1), (-1, 0)] # E, S, W, N | |
| def is_valid(r, c): | |
| return 0 <= r < rows and 0 <= c < cols and grid[r][c] != '#' | |
| def dijkstra(start_r, start_c, part2=False): | |
| q = [(0, start_r, start_c, 0)] # (cost, row, col, direction) | |
| visited = set() | |
| min_cost = float('inf') | |
| best_paths = set() | |
| while q: | |
| cost, r, c, d = heapq.heappop(q) | |
| if (r, c, d) in visited: | |
| continue | |
| visited.add((r, c, d)) | |
| if (r, c) == end: | |
| if part2: | |
| if cost == min_cost: | |
| best_paths.add((r,c)) | |
| elif cost < min_cost: | |
| min_cost = cost | |
| best_paths = set([(r,c)]) | |
| else: | |
| return cost | |
| continue | |
| if part2 and cost > min_cost: | |
| continue | |
| # Move forward | |
| nr, nc = r + directions[d][0], c + directions[d][1] | |
| if is_valid(nr, nc): | |
| heapq.heappush(q, (cost + 1, nr, nc, d)) | |
| if part2 and cost+1 <= min_cost: | |
| best_paths.add((nr, nc)) | |
| # Rotate | |
| for nd in [(d + 1) % 4, (d - 1 + 4) % 4]: # Clockwise and counterclockwise | |
| heapq.heappush(q, (cost + 1000, r, c, nd)) | |
| if part2 and cost + 1000 <= min_cost: | |
| best_paths.add((r,c)) | |
| if part2: | |
| best_paths.add(start) | |
| return len(best_paths) | |
| return float('inf') | |
| result1 = dijkstra(start[0], start[1]) | |
| print(result1) | |
| result2 = dijkstra(start[0], start[1], part2=True) | |
| print(result2) | |
| solve() |