-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathaoc202420_checkrest.py
85 lines (67 loc) · 2.57 KB
/
aoc202420_checkrest.py
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
"""AoC 20, 2024: Race Condition."""
# Standard library imports
import collections
import pathlib
import sys
def parse_data(puzzle_input):
"""Parse input."""
grid = {
(row, col): char
for row, line in enumerate(puzzle_input.split("\n"))
for col, char in enumerate(line)
}
start = next(pos for pos, char in grid.items() if char == "S")
target = next(pos for pos, char in grid.items() if char == "E")
return walk(grid, start, target)
def part1(path, min_cheat_save=100):
"""Solve part 1."""
return find_short_cheats(path, min_cheat_save)
def part2(path, min_cheat_save=100):
"""Solve part 2."""
return find_long_cheats(path, min_cheat_save, max_dist=20)
def walk(grid, start, target):
"""Find the path from start to target in the grid"""
queue = collections.deque([(start, [start])])
seen = set()
while queue:
pos, path = queue.popleft()
if pos == target:
return path
if pos in seen:
continue
seen.add(pos)
row, col = pos
for new_pos in [(row - 1, col), (row, col + 1), (row + 1, col), (row, col - 1)]:
if grid.get(new_pos, "#") != "#" and new_pos not in seen:
queue.append((new_pos, path + [new_pos]))
return []
def find_short_cheats(path, min_save):
"""Find the number of cheats walking through a single wall (2 steps)"""
positions = {pos: time for time, pos in enumerate(path)}
num_cheats = 0
for (row, col), time in positions.items():
for cheat in [(row - 2, col), (row, col + 2), (row + 2, col), (row, col - 2)]:
if positions.get(cheat, 0) - time - 2 >= min_save:
num_cheats += 1
return num_cheats
def find_long_cheats(path, min_save, max_dist=20):
"""Find the number of cheats within the given distance"""
num_cheats = 0
for time, pos in enumerate(path):
for cheat_time, cheat_pos in enumerate(
path[time + min_save :], start=time + min_save
):
dist = abs(cheat_pos[0] - pos[0]) + abs(cheat_pos[1] - pos[1])
if dist <= max_dist and cheat_time - time - dist >= min_save:
num_cheats += 1
return num_cheats
def solve(puzzle_input):
"""Solve the puzzle for the given input."""
data = parse_data(puzzle_input)
yield part1(data)
yield part2(data)
if __name__ == "__main__":
for path in sys.argv[1:]:
print(f"\n{path}:")
solutions = solve(puzzle_input=pathlib.Path(path).read_text().rstrip())
print("\n".join(str(solution) for solution in solutions))