My Advent of Code solutions in Python.
kevinyap.ca/2019/12/going-fast-in-advent-of-code/
advent-of-code
python
Add 2022/16
+140
+140
2022/day16.py
+140
2022/day16.py
···
1
+
import fileinput
2
+
from collections import deque, defaultdict
3
+
from functools import reduce
4
+
from itertools import permutations
5
+
from concurrent.futures import ProcessPoolExecutor
6
+
7
+
from utils import parse_line, memoize
8
+
9
+
from more_itertools import set_partitions
10
+
11
+
12
+
GRAPH = defaultdict(set)
13
+
FLOW = {}
14
+
15
+
# Parse input.
16
+
for line in fileinput.input():
17
+
valve, rate, tunnels = parse_line(r"Valve (\w+) has flow rate=(\d+); tunnels? leads? to valves? (.+)", line)
18
+
19
+
for tunnel in tunnels.split(", "):
20
+
GRAPH[valve].add(tunnel)
21
+
22
+
FLOW[valve] = rate
23
+
24
+
# Construct compressed graph
25
+
def bfs(start, end):
26
+
seen = set()
27
+
horizon = deque([(start, 0)])
28
+
while horizon:
29
+
curr, dist = horizon.popleft()
30
+
if curr in seen:
31
+
continue
32
+
33
+
if curr == end:
34
+
return dist
35
+
36
+
seen.add(curr)
37
+
38
+
for n in GRAPH[curr]:
39
+
horizon.append((n, dist + 1))
40
+
41
+
42
+
# Solve part 1.
43
+
START = "AA"
44
+
START_TO_REAL = {}
45
+
REAL = [v for v in GRAPH if FLOW[v] > 0]
46
+
COMPRESSED = defaultdict(set)
47
+
48
+
for valve in REAL:
49
+
START_TO_REAL[valve] = bfs(START, valve)
50
+
51
+
for a, b in permutations(REAL, 2):
52
+
dist = bfs(a, b)
53
+
COMPRESSED[a].add((b, dist))
54
+
COMPRESSED[b].add((a, dist))
55
+
56
+
57
+
@memoize
58
+
def search(nodes, max_time):
59
+
best = 0
60
+
seen = {}
61
+
horizon = deque()
62
+
63
+
# Recompute compressed graph based on limited nodes available.
64
+
compressed = {
65
+
a: [(n, d) for n, d in b if n in nodes]
66
+
for a, b in COMPRESSED.items() if a in nodes
67
+
}
68
+
69
+
# Seed start positions with time and pressure based on the time it takes
70
+
# to get from AA to that valve, + 1 minute to open the valve.
71
+
for node in nodes:
72
+
start_time = START_TO_REAL[node] + 1
73
+
start_pressure = (max_time - start_time) * FLOW[node]
74
+
state = (start_time, node, start_pressure, frozenset([node]))
75
+
horizon.append(state)
76
+
77
+
while horizon:
78
+
time, curr, pressure, opened = horizon.popleft()
79
+
key = (curr, time, opened)
80
+
81
+
if seen.get(key, -1) > pressure:
82
+
continue
83
+
84
+
seen[key] = pressure
85
+
86
+
if time >= max_time:
87
+
continue
88
+
elif len(opened) >= len(nodes):
89
+
continue
90
+
91
+
if pressure > best:
92
+
best = pressure
93
+
94
+
# move
95
+
for n, dist in compressed[curr]:
96
+
if n not in opened:
97
+
new_time = time + dist + 1
98
+
extra_pressure = (max_time - new_time) * FLOW[n]
99
+
new_opened = frozenset([n]) | opened
100
+
new_state = (new_time, n, pressure + extra_pressure, new_opened)
101
+
102
+
horizon.append(new_state)
103
+
104
+
return best
105
+
106
+
print("Part 1:", search(COMPRESSED, 30))
107
+
108
+
109
+
# Solve part 2.
110
+
def dual_search(partition):
111
+
a, b = partition
112
+
return search(frozenset(a), 26) + search(frozenset(b), 26)
113
+
114
+
115
+
part_2 = 0
116
+
partitions = list(set_partitions(REAL, 2))
117
+
118
+
# This is not the optimal solution yet. If it were, I wouldn't be parallelizing.
119
+
parallel = True
120
+
121
+
if parallel:
122
+
with ProcessPoolExecutor() as pool:
123
+
try:
124
+
results = pool.map(dual_search, partitions)
125
+
except:
126
+
pass
127
+
part_2 = reduce(max, results)
128
+
129
+
else:
130
+
# Render loading bar if tqdm installed.
131
+
try:
132
+
from tqdm import tqdm as loading
133
+
except ImportError:
134
+
loading = lambda x, total: x
135
+
136
+
for p in loading(set_partitions(REAL, 2), total=len(partitions)):
137
+
part_2 = max(part_2, dual_search(p))
138
+
139
+
print("Part 2:", part_2)
140
+