aoc-2022/16.py

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)