48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
import sys
|
|
import re
|
|
|
|
pat = re.compile('Valve ([A-Z]{2}) has flow rate=([0-9]+); tunnels? leads? to valves? ([ ,A-Z]+)')
|
|
P = dict()
|
|
for i,l in enumerate(sys.stdin.read().splitlines()):
|
|
g = pat.search(l)
|
|
p = int(g[2])
|
|
P[g[1]] = (i, p, g[3].split(", "))
|
|
|
|
STEPS = int(sys.argv[1])
|
|
ALLO = 0
|
|
for _,(i,p,_) in P.items():
|
|
if p == 0:
|
|
ALLO = ALLO | (1<<i)
|
|
|
|
cache = {}
|
|
def step(pos="AA", left=STEPS, opens=ALLO):
|
|
"""
|
|
opens store valves openess 0: closed, 1: open
|
|
1 2 3 4 5 6 7 8 9 10
|
|
eg: 0b1000011110 -> 0 1 1 1 1 0 0 0 0 1
|
|
|
|
return : score = max(p*left + open valve if not open and continue, not open valve and continue)
|
|
"""
|
|
if left <= 0:
|
|
return 0
|
|
key = (pos, left, opens)
|
|
r = cache.get(key, -1)
|
|
if r >= 0:
|
|
return r
|
|
(i, p, nh) = P[pos]
|
|
is_open = opens & (1 << i)
|
|
#print(minutes, path, i, bool(is_open), p, nh)
|
|
a = 0
|
|
if not is_open: # also means 0 - worth opening valve ?
|
|
#print("open", i)
|
|
left -= 1 # cost to open
|
|
opens = opens | (1<<i)
|
|
a = p*(left+30-STEPS) + max(step(n, left-1, opens) for n in nh)
|
|
b = max(step(n, left-1, opens) for n in nh)
|
|
r = max(a,b)
|
|
cache[key] = r
|
|
return r
|
|
|
|
print(step())
|
|
|