day 16 speed up search by shrinking the graph

day16/search.py

@@ -12,7 +12,8 @@ import sys, os; sys.path.append(os.path.join(os.path.dirname(__file__), "../lib"
 import astar
 
 def search():
-    V = sorted(v for v,_,_ in G) # vertices
+    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
@@ -35,6 +36,31 @@ def search():
     # 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:
@@ -44,10 +70,7 @@ def search():
     # (unless there is only 1 minute left)
     def heuristic(node):
         v, minutes, closed = node
-        if closed == all_open:
-            m = minutes
-        else:
-            m = min(minutes,1)
+        m = min(2, minutes)
         return m*pressure(closed)
 
     def pressure(bits):
@@ -57,9 +80,11 @@ def search():
                 pressure += r
         return pressure
 
+    assert pressure(all_open) == 0
+
     def is_goal(n):
         v, minutes, closed = n
-        return minutes == 0
+        return minutes == 0 or closed == all_open
 
     info = {}
     def neighbors(n):
@@ -71,19 +96,24 @@ def search():
         if minutes <= 0:
             pass
         elif closed == all_open:
-            c = pressure(closed) * minutes
-            yield c, (v, 0, closed)
+            yield 0, (v, 0, closed)
         else:
-            c = pressure(closed)
-            if v&closed and R[v]:
-                yield c, (v, minutes-1, closed & ~v)
-            for e in E[v]:
-                yield c, (e, minutes-1, closed)
-
+            # 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)*30 - c)
+    print(pressure(all_closed)*minutes - c)
 
     #maxpair = []
     #def pairs():