day 16 part 2 python, finally

ugh

day16/sol.py

@@ -30,13 +30,11 @@ def save():
 
 #import astar
 
-score = {'AA': (0, 0, [])}
+def part1():
+    score = {'AA': (0, 0, [])}
 
-# find the largest reward we can get in n steps
-# then find the largest reward we can get in n+1 steps
-potential = sum(r for _,r,_ in G)
-minutes = 30
-for _ in range(minutes):
+    minutes = 30
+    for _ in range(minutes):
         minutes -= 1
         next = {}
         for v, r, E in G:
@@ -75,4 +73,151 @@ for _ in range(minutes):
         print("%d minutes left" % minutes)
         for v, (r, flow, open) in sorted(score.items(), key=lambda x: x[1]):
             print("\t", v, r, "\t", ",".join(open))
-print(max(r for r,_,_ in score.values()))
+    print(max(r for r,_,_ in score.values()))
+
+def part2():
+    H = {x[0]: x for x in G}
+    def actions(v, open):
+        _, r, E = H[v]
+        if v in open:
+            yield (v, 0, [])
+        else:
+            yield (v, 0, [])
+            yield (v+'o', r, [v])
+        # move to somewhere else
+        for e in E:
+            if e in open:
+                e = e+'o'
+            yield (e, 0, [])
+
+    score = {('AA', 'AA'): (0, 0, [])}
+    minutes = 26
+    for _ in range(minutes):
+        next = {}
+        minutes -= 1
+
+        for (me, elephant) in score:
+            reward, flow, open = score[(me, elephant)]
+            for v, r, o in actions(me.strip('o'), open):
+                for v1, r1, o1 in actions(elephant.strip('o'), open+o):
+                    reward_ = reward + (r+r1)*minutes
+                    flow_ = flow + r + r1
+                    open_ = open + o + o1
+                next.setdefault((v, v1), []).append((reward_, flow_, open_))
+
+        for k in next:
+            next[k] = max(next[k])
+
+        score = next
+        print("%d minutes left" % minutes)
+        for v, (r, flow, open) in sorted(score.items(), key=lambda x: x[1]):
+            print("\t", v, r, "\t", ",".join(open))
+    print(max(r for r,_,_ in score.values()))
+
+#part2()
+
+def part2_b():
+    score = {'AA': (0, 0, [])}
+
+    minutes = 26
+    for _ in range(minutes):
+        minutes -= 1
+        next = {}
+        for v, r, E in G:
+            vo = v + 'o'
+            s = []
+            o = []
+            if vo in score:
+                reward, flow, open = score[vo]
+                o.append((reward, flow, open))
+            if v in score:
+                reward, flow, open = score[v]
+                # stay, don't open valve
+                s.append((reward, flow, open))
+                # stay in place, open valve
+                if v not in open:
+                    reward += r*minutes
+                    o.append((reward, flow+r, open+[v]))
+            # move here from somewhere else
+            for e in E:
+                if e in score:
+                    if v in score[e][-1]:
+                        o.append(score[e])
+                    else:
+                        s.append(score[e])
+                eo = e+'o'
+                if eo in score:
+                    if v in score[eo][-1]:
+                        o.append(score[eo])
+                    else:
+                        s.append(score[eo])
+            if s:
+                next[v] = max(s)
+            if o:
+                next[vo] = max(o)
+        score = next
+        print("%d minutes left" % minutes)
+        for v, (r, flow, open) in sorted(score.items(), key=lambda x: x[1]):
+            print("\t", v, r, "\t", ",".join(open))
+
+    maxpair = []
+    for r,_,open in score.values():
+        o = frozenset(open)
+        for s,_,open in score.values():
+            if o.isdisjoint(open):
+                maxpair.append(r+s)
+    print(max(maxpair))
+
+#part2_b()
+
+def part2_c():
+    V = sorted(v for v,_,_ in G)
+    B = {v: 1<<i for i, v in enumerate(V)}
+    E = {B[v]: [B[e] for e in edges] for v,_,edges in G}
+    R = {B[v]: r for v,r,_ in G}
+
+    all_open = sum(B.values())
+
+    aa = B['AA']
+    states = [(aa, 0, 0)]
+    minutes = 26
+    for _ in range(minutes):
+        print(minutes, len(states))
+        minutes -= 1
+        next = []
+        seen = set()
+        states.sort(key=lambda x: -x[2])
+        for v, open, reward in states:
+            if (v,open) in seen:
+                continue
+            seen.add((v,open))
+            #next.append((v, open, reward))
+            if open == all_open:
+                next.append((v, open, reward))
+            else:
+                #next.append((v, open, reward))
+                r = R[v]
+                if r and not v&open:
+                    next.append((v, open|v, reward+r*minutes))
+                for e in E[v]:
+                    next.append((e, open, reward))
+        states = next
+
+    print(max(r for _,_,r in states))
+
+    best = {}
+    for _,o,r in states:
+        if o in best and best[o] >= r:
+            continue
+        best[o] = r
+
+    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()))
+
+part2_c()