day 16 python part 2 solve with A*, and cleanup

day16/search.py

@@ -11,18 +11,27 @@ for line in open("input"):
 import sys, os; sys.path.append(os.path.join(os.path.dirname(__file__), "../lib"))
 import astar
 
-def search():
+def solve():
-    G2 = sorted(G, key=lambda x: (-x[1],x[0]))
+    G.sort(key=lambda x: (-x[1],x[0]))
-    V = [v for v,_,_ in G2] # vertices
+    B = {v: 1<<i for i,(v,_,_) in enumerate(G)} # bitmasks
-    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} # E[b] = edges of b
-    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 if r} # R[b] -> rate
-    R = {B[v]: r for v,r,_ in G} # rewards: R[b] = reward
 
-    all_closed = sum(B.values())
+    AA = B['AA']
+
+    all_closed = sum(R.keys())
     all_open = 0
 
-    minutes = 30
+    # TODO: memoize this
-    start = (B['AA'], minutes, all_closed)
+    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
@@ -33,9 +42,13 @@ def search():
     # 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
+    # every pair of rooms (floyd-warshall), and then use that to build a
-    # paths from the starting room to rooms with a valve
+    # weighted graph which only has paths from any room to rooms with a valve.
-    # and from any room with a valve to any other room 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
@@ -48,47 +61,47 @@ def search():
                 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
+    W = {} # weighted edges
-    for u in W:
+    for u in E:
-        for v in W:
+        W[u] = []
+        for v in R:
             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])
+    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
-        m = min(2, minutes)
+        # assume we can open a valve every 2 minutes
-        return m*pressure(closed)
+        # 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 pressure(bits):
+    def is_goal(node):
-        pressure = 0
+        v, minutes, closed = node
-        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):
+    def neighbors(node):
-        v, minutes, closed = n
+        v, minutes, closed = node
         if minutes not in info:
             print(info)
             info[minutes] = 0
@@ -103,7 +116,7 @@ def search():
             can_move = False
             for e, dist in W[v]:
                 t = dist + 1
-                if e&closed and t <= minutes and R[e]:
+                if e&closed and t <= minutes:
                     can_move = True
                     c = pressure(closed)*t
                     yield c, (e, minutes-t, closed&~e)
@@ -111,17 +124,93 @@ def search():
             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(pressure(all_closed)*minutes - c)
+    print(max_pressure*minutes - c)
 
-    #maxpair = []
+    def heuristic2(node):
-    #def pairs():
+        if is_goal2(node):
-    #    O = sorted(best.keys())
+            return 0
-    #    for i in range(len(O)):
+        v1, v2, min1, min2, closed = node
-    #        for j in range(i,len(O)):
+        # if the players are out of sync, assume the other player
-    #            if not O[i] & O[j]:
+        # will close one valve when they catch up
-    #                yield(best[O[i]]+best[O[j]])
+        if min1 != min2:
-    #print(max(pairs()))
+            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
 
-search()
+    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()