from queue import PriorityQueue lines = [line.split() for line in open("16.input")] valves = {line[1]: (int(line[4].split("=")[1][:-1]), [n.replace(",", "") for n in line[9:]]) for line in lines} useful_valves = set(k for k, v in valves.items() if v[0] > 0) # Find all relevant distances dists = {} for a in list(useful_valves) + ["AA"]: q = PriorityQueue() dists[(a, a)] = 0 q.put((0, a)) while not q.empty(): _, current = q.get() dist = dists[(a, current)] _, nbrs = valves[current] for n in nbrs: old_dist = dists.get((a, n), 100000) if old_dist > dist + 1: dists[(a, n)] = dist + 1 q.put((dists[(a, n)], n)) def search(currents, opened, remainings, score, paths): old_score = score_cache.get((currents, remainings), -1) if score > old_score: score_cache[(currents, remainings)] = score else: return score nexts = useful_valves - opened scores = [score] _, i = max((rem, i) for i, rem in enumerate(remainings)) current = currents[i] remaining = remainings[i] remainings = list(remainings) currents = list(currents) for n in nexts: d = dists[(current, n)] if d >= remaining: continue flow, _ = valves[n] remainings[i] = remaining - d - 1 currents[i] = n new_score = score + (remaining - d - 1) * flow scores.append(search(tuple(currents), opened.union(set([n])), tuple(remainings), new_score, [])) return max(scores) score_cache = {} score = search(('AA', ), set(), (30, ), 0, []) print(score) score_cache = {} score = search(('AA', 'AA'), set(), (26, 26), 0, ([], [])) print(score)