/
aoc16e.py
117 lines (106 loc) · 3.86 KB
/
aoc16e.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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
from aoc_util import get_data_lines
import re
import heapq
from typing import Dict, Tuple
import itertools
import numpy as np
data = get_data_lines(16)
valvemap = {}
valverates = {}
for line in data:
m = re.match(
r"Valve (\S+) has flow rate=(\d+); tunnels? leads? to valves? (.*)", line
)
if m is None:
print("ERR: " + line)
name = m.group(1)
rate = int(m.group(2))
rest = [dest.strip() for dest in m.group(3).split(",")]
valvemap[name] = tuple(rest)
valverates[name] = rate
shortest: Dict[Tuple[str, str], int] = {}
for (waypoint, nexts) in valvemap.items():
shortest[(waypoint, waypoint)] = 0
for dest in nexts:
shortest[(waypoint, dest)] = 1
for waypoint in valvemap:
for first in valvemap:
for last in valvemap:
if first != last:
shortest[(first, last)] = min(
shortest.get((first, last), 999),
shortest.get((first, waypoint), 999)
+ shortest.get((waypoint, last), 999),
)
to_int = dict(
zip([v for v in valverates if valverates[v] > 0] + ["AA"], range(len(valvemap)))
)
aa_int = to_int["AA"]
distmat = np.ndarray((len(to_int), len(to_int)), dtype=int)
ratevec = np.ndarray((len(to_int),), dtype=int)
for (froms, fromi) in to_int.items():
ratevec[fromi] = valverates[froms]
for (tos, toi) in to_int.items():
distmat[fromi, toi] = shortest[(froms, tos)]
distmat[fromi, fromi] = 1
distmat[fromi, aa_int] = 999
ratelookup = np.zeros((1,), dtype=int)
for idx in range(aa_int):
ratelookup = np.append(ratelookup, ratelookup + ratevec[idx])
sum_of_valverates = ratevec.sum()
# Things in all_spots are: (sort key, pressure released so far, location, open valves, time, old_time)
all_spots = [(0, 0, aa_int, 0, 0, 0)]
best_so_far = -1
best_to_26 = {}
been_here = {}
sample = 0
max_len_all_spots = 0
while all_spots:
(_, released_so_far, loc, open_valves, time, otime) = heapq.heappop(all_spots)
sample += 1
max_len_all_spots = max(max_len_all_spots, len(all_spots))
if been_here.get((loc, open_valves, time), -1) >= released_so_far:
continue
been_here[(loc, open_valves, time)] = max(
released_so_far, been_here.get((loc, open_valves, time), -1)
)
been_here[(loc, open_valves, time + 1)] = max(
released_so_far, been_here.get((loc, open_valves, time), -1)
)
if otime < 26 and time >= 26:
at_26 = released_so_far - ratelookup[open_valves] * (time - 26)
best_to_26[open_valves] = max(best_to_26.get(open_valves, -1), at_26)
if time == 30:
best_so_far = max(best_so_far, released_so_far)
continue
nreleased = ratelookup[open_valves] + released_so_far
for (dest, desttime) in enumerate(distmat[loc]):
ntime = time + desttime
if ntime <= 30:
nreleased = ratelookup[open_valves] * desttime + released_so_far
npot = nreleased + (30 - ntime) * ratelookup[open_valves]
heapq.heappush(
all_spots, (-npot, nreleased, dest, open_valves, ntime, time)
)
if loc != aa_int:
valve = 1 << loc
if not (valve & open_valves):
ntime = time + 1
nreleased = ratelookup[open_valves] + released_so_far
nopen_valves = open_valves | valve
npot = nreleased + (30 - ntime) * ratelookup[nopen_valves]
heapq.heappush(
all_spots, (-npot, nreleased, loc, nopen_valves, ntime, time)
)
# print("diag: ", len(all_spots), time)
# part 1
print(best_so_far)
# print(
# f"Samples nreleased + (30 - ntime) * ratelookup[open_valves]: {sample}, max_len: {max_len_all_spots}"
# )
# part 2
best_so_far = -1
for (op1, op2) in itertools.combinations(best_to_26, 2):
if not (op1 & op2):
best_so_far = max(best_so_far, best_to_26[op1] + best_to_26[op2])
print(best_so_far)