-
Notifications
You must be signed in to change notification settings - Fork 0
/
Day16.py
64 lines (56 loc) · 1.97 KB
/
Day16.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
import re
from itertools import combinations
from tqdm import tqdm
g = {}
rates = {}
with open('inputs/day16.txt') as f:
lines = f.read().strip().split('\n')
for line in lines:
tokens = line.split(' ')
cur = tokens[1]
g[cur] = tuple(map(lambda x: x.rstrip(','), tokens[9:]))
rates[cur] = int(tokens[4][5:-1])
# print(g)
def dfs(source):
distance = {}
def aux(node, travelled = 0):
neighbors = g[node]
for neighbor in neighbors:
if neighbor != source:
if neighbor in distance:
if travelled + 1 < distance[neighbor]:
distance[neighbor] = travelled + 1
aux(neighbor, travelled + 1)
else:
distance[neighbor] = travelled + 1
aux(neighbor, travelled + 1)
aux(source)
return distance
paths = {}
for source in g:
paths[source] = dfs(source)
def visit(node, unvisited, remaining, pressure):
if remaining < 0:
return None
if len(unvisited) == 0:
return pressure
max_pressure = pressure
for target in unvisited:
distance = paths[node][target]
time = distance + 1
new_pressure = visit(target, unvisited - {target}, remaining - time, pressure + rates[target] * (remaining - time))
if new_pressure is not None and new_pressure > max_pressure:
max_pressure = new_pressure
return max_pressure
valves = {node for (node, rate) in rates.items() if rate > 0 and node != 'AA'}
print(visit('AA', valves, 30, 0))
max_pressure = 0
# Only need to test half as the remaining are just inverse
for i in tqdm(range(1, len(valves)//2 + 1)):
for me in tqdm(combinations(valves, i)):
me = set(me)
elephant = valves - me
me_pressure = visit('AA', me, 26, 0)
elephant_pressure = visit('AA', elephant, 26, 0)
max_pressure = max(max_pressure, me_pressure + elephant_pressure)
print(max_pressure)