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, defaultdict
3
4from utils import Point, N, S, E, W, DIRS
5
6
7SLOPES = {
8 '^': S,
9 'v': N,
10 '<': W,
11 '>': E,
12}
13
14
15def get_neighbours(graph, node, part_2=False):
16 """In part 2, we ignore slopes."""
17 if graph.get(node) == '#':
18 return
19
20 if not part_2:
21 if graph.get(node) in SLOPES:
22 yield node + SLOPES[graph.get(node)]
23 return
24
25 for d in DIRS:
26 np = node + d
27 if np not in graph:
28 continue
29
30 neighbour = graph.get(np)
31
32 if neighbour == '#':
33 continue
34
35 if not part_2 and neighbour in SLOPES and d != SLOPES[neighbour]:
36 continue
37
38 yield np
39
40
41def longest_path(graph, start, end, part_2=False):
42 BITMASK_MAPPING = {k: i for i, k in enumerate(graph)}
43 horizon = [(start, 0, 0b0)]
44 best = 0
45
46 while horizon:
47 curr, dist, seen = horizon.pop()
48
49 if curr == end:
50 best = max(best, dist)
51 continue
52
53 mask = (1 << BITMASK_MAPPING[curr])
54
55 if seen & mask:
56 continue
57
58 new_seen = seen | mask
59
60 for neighbour, weight in graph[curr]:
61 horizon.append((neighbour, dist + weight, new_seen))
62
63 return best
64
65
66def compress_graph(graph, part_2=False):
67 # Sort all points in graph by their degree.
68 degrees = defaultdict(set)
69 for node in graph:
70 degrees[len(list(get_neighbours(graph, node, part_2)))].add(node)
71
72 key_points = degrees[1] | degrees[3] | degrees[4]
73
74
75 # Find the distance from node to all other "key points" it can reach.
76 def bfs(start):
77 horizon = deque([(start, 0)])
78 seen = set()
79
80 while horizon:
81 curr, dist = horizon.pop()
82
83 if curr != start and curr in key_points:
84 yield curr, dist
85 continue
86
87 if curr in seen:
88 continue
89
90 seen.add(curr)
91
92 for neighbour in get_neighbours(graph, curr, part_2):
93 horizon.appendleft((neighbour, dist + 1))
94
95
96 # Create the compressed weighted graph.
97 compressed = defaultdict(list)
98
99 for node in key_points:
100 for neighbour, weight in bfs(node):
101 compressed[node].append((neighbour, weight))
102
103 return compressed
104
105
106# Parse problem input.
107GRAPH = {}
108START = None
109END = None
110for y, line in enumerate(fileinput.input()):
111 for x, c in enumerate(line.strip()):
112 p = Point(x, y)
113 if c == '.':
114 if y == 0:
115 START = p
116
117 END = p
118
119 GRAPH[p] = c
120
121print("Part 1:", longest_path(compress_graph(GRAPH), START, END))
122print("Part 2:", longest_path(compress_graph(GRAPH, part_2=True), START, END, part_2=True))
123