problem16.py

from itertools import product
fh = open("input16")
G = {}
Flow = {}
for line in fh:
    line = line.rstrip()
    first, second = line.split(";")
    valve = first[6:8]
    Flow[valve] = int(first.split("=")[1])
    l = second.split(", ")
    l[0] = l[0][-2:]
    G[valve] = l

def solve(t, position, valves, score, cache):
    if t == 1:
        return score
    if (t, position, valves) in cache:
        return score + cache[(t, position, valves)]
    strategies = [solve(t-1, p, valves, score, cache) for p in G[position]]
    r = max(strategies)
    if Flow[position] > 0 and position not in valves:
        alt_strat = solve(t-1, position, valves | frozenset([position]), score + (t-1) * Flow[position], cache)
        r = max(alt_strat, r)
    cache[(t, position, valves)] = r - score
    return r

def solve2(t, pos1, pos2, valves, score, cache):
    if len(valves) == 15:
        return score
    if t == 1:
        return score
    if t == 2:
        if (Flow[pos1] > 0 and pos1 not in valves):
            score += Flow[pos1]
        if (Flow[pos2] > 0 and pos2 not in valves):
            score += Flow[pos2]
        return score
    if (t, pos1, pos2, valves) in cache:
        return score + cache[(t, pos1, pos2, valves)]
    if (t, pos2, pos1, valves) in cache:
        return score + cache[(t, pos2, pos1, valves)]
    strategies = [solve2(t-1, p1, p2, valves, score, cache) for p1, p2 in product(G[pos1], G[pos2])]
    if Flow[pos1] > 0 and pos1 not in valves:
        strategies.extend([solve2(t-1, pos1, p2, valves | frozenset([pos1]), score + (t-1) * Flow[pos1], cache) for p2 in G[pos2]])
    if Flow[pos2] > 0 and pos2 not in valves:
        strategies.extend([solve2(t-1, p1, pos2, valves | frozenset([pos2]), score + (t-1) * Flow[pos2], cache) for p1 in G[pos1]])
    if pos1 != pos2 and Flow[pos1] > 0 and Flow[pos2] > 0 and pos1 not in valves and pos2 not in valves:
        strategies.append(solve2(t-1, pos1, pos2, valves | frozenset([pos1, pos2]), score + (t-1) * (Flow[pos1] + Flow[pos2]), cache))
    r = max(strategies)
    cache[(t, pos1, pos2, valves)] = r - score
    return r

cache = {}
valves = frozenset()
pomme = solve2(26, "AA", "AA", valves, 0, cache)