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 solve():
    G.sort(key=lambda x: (-x[1],x[0]))
    B = {v: 1<<i for i,(v,_,_) in enumerate(G)} # bitmasks
    E = {B[v]: [B[e] for e in edges] for v,_,edges in G} # E[b] = edges of b
    R = {B[v]: r for v,r,_ in G if r} # R[b] -> rate

    AA = B['AA']

    all_closed = sum(R.keys())
    all_open = 0

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

    assert pressure(all_open) == 0
    max_pressure = pressure(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 use that to build a
    # weighted graph which only has paths from any room to rooms with a valve.
    #
    # not only does this make our search space smaller,
    # it also helps by making it so that the cost changes on every step
    # (since opening a valve is the only thing that actually changes the pressure)
    # giving A* a much clearer signal about which paths are worth exploring
    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 = {} # weighted edges
    for u in E:
        W[u] = []
        for v in R:
            if (u,v) in dist:
                W[u].append((v, dist[u,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)
    #
    # we can keep doing that until there are no valves left to open
    # or there is no time left.
    #
    # note that the nodes with the largest flow rate are assigned
    # the lowest position in the bitmap, so clearing the bits from
    # low to high will always give us the optimal order
    def heuristic(node):
        v, minutes, closed = node
        # assume we can open a valve every 2 minutes
        # how much would that cost?
        c = 0
        while closed and minutes > 0:
            c += pressure(closed) * min(minutes,2)
            closed &= (closed-1)
            minutes -= 2
        return c

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

    info = {}
    def neighbors(node):
        v, minutes, closed = node
        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:
                    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)

    minutes = 30
    start = (AA, minutes, all_closed)

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

    def heuristic2(node):
        if is_goal2(node):
            return 0
        v1, v2, min1, min2, closed = node
        # if the players are out of sync, assume the other player
        # will close one valve when they catch up
        if min1 != min2:
            closed &= closed - 1
        # assume we can open a valve every minute remaining
        # how much would that cost?
        c = 0
        pr = pressure(closed)
        m = min(min1, min2)
        while closed and m > 0:
            c += pr
            tmp = closed
            closed &= (closed-1)
            pr -= R[tmp-closed]
            m -= 1
        return c

    def is_goal2(node):
        _, _, min1, min2, closed = node
        return min1 == 0 and min2 == 0 or closed == all_open

    def neighbors2(node):
        v1, v2, min1, min2, closed = node

        if min(min1,min2) not in info:
            print(info)
        info.setdefault(min1, 0)
        info.setdefault(min2, 0)
        info[min1] += 1
        info[min2] += 1

        if min1 <= 0 and min2 <= 0:
            pass
        elif closed == all_open:
            yield 0, (v1, v2, 0, 0, closed)
        else:
            moved = False
            # either player can move
            # but we can't open a valve that would take less time to open
            # than the other player has already used. we've already paid
            # the cost for _not_ opening those valves and we can't
            # retroactively change that (no time travel)

            # if both players are in the same spot and have the same amount of
            # time remaining, then only let one of them move in order to break
            # symmetries (this can only happen in the start state)

            # TODO: are there more symmetries we can break?

            # move to a closed valve and open it
            discount1 = max(min1-min2, 0)
            for e, dist in W[v1]:
                t = dist + 1
                if e&closed and discount1 <= t <= min1:
                    moved = True
                    c = pressure(closed)*(t-discount1)
                    yield c, (e, v2, min1-t, min2, closed&~e)
            if (v1, min1) != (v2, min2):
                discount2 = max(min2-min1, 0)
                for e, dist in W[v2]:
                    t = dist + 1
                    if e&closed and discount2 <= t <= min2:
                        moved = True
                        c = pressure(closed)*(t-discount2)
                        yield c, (v1, e, min1, min2-t, closed&~e)

            # are there no moves left?
            # then wait out the timer
            if not moved:
                yield pressure(closed)*min(min1,min2), (v1, v2, 0, 0, closed)

    minutes = 26
    start2 = (AA, AA, minutes, minutes, all_closed)
    info.clear()
    c, _, path = astar.search(start2, is_goal2, neighbors2, heuristic2)
    print(c)
    print(max_pressure*minutes - c)

solve()