My Advent of Code solutions in Python.
kevinyap.ca/2019/12/going-fast-in-advent-of-code/
advent-of-code
python
1import fileinput
2import heapq
3
4from utils import Point, DIRS
5
6
7def minimize_heat_loss(graph, start, goal=None, part_2=False):
8 def gen_neighbours(node, last_dir, last_count):
9 for d in DIRS:
10 n = node + d
11 if n not in graph:
12 continue
13
14 if last_dir is not None:
15 # No U-turns allowed.
16 if last_dir == -d:
17 continue
18
19 if part_2:
20 # Early ultra crucible turn.
21 if last_count < 4 and d != last_dir:
22 continue
23
24 # Wobbly ultra crucible.
25 if last_count == 10 and d == last_dir:
26 continue
27 else:
28 # Unstable regular crucible.
29 if last_count == 3 and last_dir == d:
30 continue
31
32 yield n, graph[n], d
33
34 horizon = [(0, start, None, 0)]
35 seen = set()
36
37 while horizon:
38 depth, curr, last_dir, last_count = heapq.heappop(horizon)
39
40 if (curr, last_dir, last_count) in seen:
41 continue
42
43 seen.add((curr, last_dir, last_count))
44
45 # Check if at the bottom-right, and that our last 4 moves
46 # were in the same direction if we're solving part 2.
47 if curr == goal and (not part_2 or last_count >= 4):
48 return depth
49
50 for neighbour, weight, nd in gen_neighbours(curr, last_dir, last_count):
51 new_cost = weight + depth
52 new_count = 1 if nd != last_dir else last_count + 1
53 heapq.heappush(horizon, (new_cost, neighbour, nd, new_count))
54
55
56# Read problem input.
57graph = {}
58for y, line in enumerate(fileinput.input()):
59 for x, c in enumerate(line.strip()):
60 graph[Point(x, y)] = int(c)
61
62start = Point(0, 0)
63end = Point(x, y)
64
65# Solve problem.
66print("Part 1:", minimize_heat_loss(graph, start, end, part_2=False))
67print("Part 2:", minimize_heat_loss(graph, start, end, part_2=True))