51 lines
1.6 KiB
Python
51 lines
1.6 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(pos1="AA", pos2="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 = (min(pos1, pos2), max(pos1,pos2), left, opens)
|
|
r = cache.get(key, -1)
|
|
if r >= 0:
|
|
return r
|
|
(i1, p1, nh1) = P[pos1]
|
|
(i2, p2, nh2) = P[pos2]
|
|
isopen1 = opens & (1 << i1)
|
|
isopen2 = opens & (1 << i2)
|
|
a = b = c = d = 0
|
|
if pos1 != pos2 and not isopen1 and not isopen2: # open both if not same, but don't move
|
|
a = p1*(left+26-STEPS-1) + p2*(left+26-STEPS-1) + step(pos1, pos2, left-1, opens|(1<<i1)|(1<<i2))
|
|
if not isopen1: # p1 opens, p2 moves
|
|
b = p1*(left+26-STEPS-1) + max(step(pos1, n2, left-1, opens|(1<<i1)) for n2 in nh2)
|
|
if not isopen2: # p1 moves, p2 opens
|
|
c = p2*(left+26-STEPS-1) + max(step(n1, pos2, left-1, opens|(1<<i2)) for n1 in nh1)
|
|
# both moves, no new open
|
|
d = max(step(n1, n2, left-1, opens) for n1 in nh1 for n2 in nh2)
|
|
r = max(a,b,c,d)
|
|
cache[key] = r
|
|
return r
|
|
|
|
print(step())
|
|
print(len(cache))
|