My Advent of Code solutions in Python.
kevinyap.ca/2019/12/going-fast-in-advent-of-code/
advent-of-code
python
1import fileinput
2from collections import deque
3
4from utils import Point, N, E, S, W
5
6
7def search(graph, start, goal):
8 def gen_neighbours(node, depth):
9 """Returns all valid neighbours for a given node."""
10 for n in node.neighbours_4():
11 if graph.get(n, ' ') == '.':
12 if n not in get_blizzard_locs(depth + 1):
13 yield n
14
15 if node not in get_blizzard_locs(depth + 1):
16 yield node
17
18 part_1 = None
19 part_2 = None
20
21 horizon = deque([(0, 0, start)]) # priority, checkpoint, node
22 seen = set()
23
24 while horizon:
25 depth, checkpoint, curr = horizon.popleft()
26
27 if (depth, checkpoint, curr) in seen:
28 continue
29
30 seen.add((depth, checkpoint, curr))
31
32 # Reached the goal for the second time after making it back to the start.
33 if curr == goal and checkpoint == 2:
34 part_2 = depth
35 print("Part 2:", part_2)
36 return part_1, part_2
37
38 # Reached the goal for the first time.
39 elif curr == goal and checkpoint == 0:
40 if part_1 is None:
41 part_1 = depth
42 print("Part 1:", part_1)
43 checkpoint = 1
44
45 # Reached the start after touching the goal the first time.
46 elif curr == start and checkpoint == 1:
47 checkpoint = 2
48
49 for n in gen_neighbours(curr, depth):
50 horizon.append((depth + 1, checkpoint, n))
51
52
53# N and S are swapped because of positive-Y direction convention.
54MAPPING = {
55 'v': N,
56 '>': E,
57 '^': S,
58 '<': W,
59}
60
61BLIZZARDS = {}
62BLIZZARD_LOCS = {}
63
64def get_blizzard_locs(step):
65 """Returns the set of all blizzard positions for the given time step."""
66 if step in BLIZZARDS:
67 return BLIZZARD_LOCS[step]
68
69 last_blizzards = BLIZZARDS[step - 1]
70 new_blizzards = []
71 for p, b in last_blizzards:
72 np = p + MAPPING[b]
73 if BOARD.get(np) == '#':
74 np = p
75 while BOARD.get(np) != '#':
76 np -= MAPPING[b]
77 np += MAPPING[b]
78
79 new_blizzards.append((np, b))
80
81 BLIZZARDS[step] = new_blizzards
82 BLIZZARD_LOCS[step] = set(p for p, _ in new_blizzards)
83 return BLIZZARD_LOCS[step]
84
85
86# Read problem input.
87BOARD = {}
88BLIZZARDS = {0: []}
89START = None
90GOAL = None
91for y, line in enumerate(fileinput.input()):
92 for x, c in enumerate(line.strip()):
93 p = Point(x, y)
94 if c == '.':
95 if y == 0:
96 START = p
97 GOAL = p
98 BOARD[p] = '.'
99 elif c in MAPPING:
100 BLIZZARDS[0].append((p, c))
101 BOARD[p] = '.'
102 else:
103 BOARD[p] = '#'
104
105# Solve problem.
106search(BOARD, START, GOAL)