adventofcode2022/day16/search.py

G = []
for line in open("input"):
    words = line.split()
    valve = words[1]
    rate = int(''.join(x for x in words[4] if x.isdigit()))
    edges = [x.strip(", ") for x in words[9:]]
    G.append((valve, rate, edges))

#print(G)

import sys, os; sys.path.append(os.path.join(os.path.dirname(__file__), "../lib"))
import astar

def search():
    G2 = sorted(G, key=lambda x: (-x[1],x[0]))
    V = [v for v,_,_ in G2] # vertices
    B = {v: 1<<i for i, v in enumerate(V)} # bitmasks
    E = {B[v]: [B[e] for e in edges] for v,_,edges in G} # edge[b] -> b
    R = {B[v]: r for v,r,_ in G} # rewards: R[b] = reward

    all_closed = sum(B.values())
    all_open = 0

    minutes = 30
    start = (B['AA'], minutes, all_closed)

    # A* search minimizes costs
    # it can't maxmize anything
    # so we'll borrow an idea from https://github.com/morgoth1145/advent-of-code/blob/2bf7c157e37b3e0a65deedc6c88e42297d813d1d/2022/16/solution.py
    # and instead say that the cost of moving from one node to the next
    # is equal to the potential pressure we could have released from the closed pipes
    # or, in other words, we'll keep track of how much pressure builds up in
    # closed pipes instead of how much pressure is released from open pipes

    # let's shink the graph by finding the shortest path between
    # every pair of rooms (floyd-warshall), and then building a graph which only has
    # paths from the starting room to rooms with a valve
    # and from any room with a valve to any other room with a valve
    dist = {}
    for v in E:
        dist[v,v] = 0
        for e in E[v]:
            dist[v,e] = 1
            dist[e,v] = 1
    for t in E:
        for u in E:
            for v in E:
                if (u,t) in dist and (t,v) in dist:
                    dist[u,v] = min(dist.get((u,v),999999), dist[u,t] + dist[t,v])

    W = {v:[] for v in R if R[v]} # weighted edges
    for u in W:
        for v in W:
            if (u,v) in dist:
                W[u].append((v, dist[u,v]))
    aa = B['AA']
    if aa not in W:
        W[aa] = []
        for v in W:
            if v != aa and (aa,v) in dist:
                W[aa].append((v, dist[aa,v]))

    print(W[aa])

    # our heuristic cost has to be <= the actual cost of getting to the goal
    # here's a simple one:
    # we know it takes at least 1 minute to open a valve,
    # and at least another minute to walk to the valve
    # so we can assign at least cost 2*pressure(closed_valves) to this node
    # (unless there is only 1 minute left)
    def heuristic(node):
        v, minutes, closed = node
        m = min(2, minutes)
        return m*pressure(closed)

    def pressure(bits):
        pressure = 0
        for v,r in R.items():
            if bits&v:
                pressure += r
        return pressure

    assert pressure(all_open) == 0

    def is_goal(n):
        v, minutes, closed = n
        return minutes == 0 or closed == all_open

    info = {}
    def neighbors(n):
        v, minutes, closed = n
        if minutes not in info:
            print(info)
            info[minutes] = 0
        info[minutes] += 1
        if minutes <= 0:
            pass
        elif closed == all_open:
            yield 0, (v, 0, closed)
        else:
            # move to a closed valve (or maybe stay in
            # the same spot) and open it
            can_move = False
            for e, dist in W[v]:
                t = dist + 1
                if e&closed and t <= minutes and R[e]:
                    can_move = True
                    c = pressure(closed)*t
                    yield c, (e, minutes-t, closed&~e)
            # wait til the end
            if can_move == False:
                yield pressure(closed)*minutes, (v, 0, closed)

    c, _, path = astar.search(start, is_goal, neighbors, heuristic)
    print(c)
    print(pressure(all_closed)*minutes - c)

    #maxpair = []
    #def pairs():
    #    O = sorted(best.keys())
    #    for i in range(len(O)):
    #        for j in range(i,len(O)):
    #            if not O[i] & O[j]:
    #                yield(best[O[i]]+best[O[j]])
    #print(max(pairs()))

search()