Backup local 2022 day 16.

2022/day16.py

@@ -6,10 +6,11 @@ import heapq
 import itertools
 import re
 import sys
+import time as time_p
 from collections import defaultdict
 from typing import FrozenSet, NamedTuple
 
-from tqdm import tqdm
+from tqdm import tqdm, trange
 
 
 class Pipe(NamedTuple):
@@ -106,21 +107,128 @@ def part_1(
 def part_2(
     start_pipe: Pipe,
     max_time: int,
-    distances: dict[tuple[Pipe, Pipe], int],
+    pipes: dict[str, Pipe],
     relevant_pipes: FrozenSet[Pipe],
+    distances: dict[tuple[Pipe, Pipe], int],
 ):
-    def compute(pipes_for_me: FrozenSet[Pipe]) -> int:
+
-        return part_1(start_pipe, max_time, distances, pipes_for_me) + part_1(
+    node_at_times: dict[
-            start_pipe, max_time, distances, relevant_pipes - pipes_for_me
+        int, dict[tuple[Pipe, Pipe], dict[FrozenSet[Pipe], int]]
+    ] = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: 0)))
+    node_at_times[0] = {(start_pipe, start_pipe): {frozenset(): 0}}
+
+    # map node + distance to
+    d1, d2, d3, d4 = 0, 0, 0, 0
+    best_flow = 0
+
+    for time in range(max_time):
+        print(
+            f"{time + 1:2d}/{max_time} - {best_flow:4d} - "
+            f"{sum(map(len, node_at_times[time].values())):7d} - "
+            f"{d1:.3f} {d2:.3f} {d3:.3f} {d4:.3f}"
         )
 
-    combs = [
+        d1, d2, d3, d4 = 0, 0, 0, 0
-        frozenset(relevant_pipes_1)
+        for (c_pipe, e_pipe), nodes in node_at_times[time].items():
-        for r in range(2, len(relevant_pipes) // 2 + 1)
+            for flowing, flow in nodes.items():
-        for relevant_pipes_1 in itertools.combinations(relevant_pipes, r)
-    ]
 
-    return max(compute(comb) for comb in tqdm(combs))
+                t1 = time_p.time()
+
+                c_best_flow = (
+                    flow
+                    + sum(pipe.flow for pipe in flowing) * (max_time - time)
+                    + sum(
+                        (
+                            pipe.flow
+                            * (
+                                max_time
+                                - time
+                                - 1
+                                - min(distances[c_pipe, pipe], distances[e_pipe, pipe])
+                            )
+                            for pipe in relevant_pipes
+                            if pipe not in flowing
+                        ),
+                        start=0,
+                    )
+                )
+
+                d1 += time_p.time() - t1
+
+                if c_best_flow < best_flow:
+                    continue
+
+                best_flow = max(
+                    best_flow,
+                    flow + sum(pipe.flow for pipe in flowing) * (max_time - time),
+                )
+
+                t1 = time_p.time()
+
+                if flowing != relevant_pipes:
+                    for c_next_s, e_next_s in itertools.product(
+                        c_pipe.tunnels, e_pipe.tunnels
+                    ):
+
+                        c_next = pipes[c_next_s]
+                        e_next = pipes[e_next_s]
+                        update_with_better(
+                            node_at_times[time + 1][c_next, e_next],
+                            flow + sum(pipe.flow for pipe in flowing),
+                            flowing,
+                        )
+
+                d2 += time_p.time() - t1
+
+                t1 = time_p.time()
+
+                if c_pipe in relevant_pipes and c_pipe not in flowing:
+                    for e_next_s in e_pipe.tunnels:
+
+                        e_next = pipes[e_next_s]
+
+                        update_with_better(
+                            node_at_times[time + 1][c_pipe, e_next],
+                            flow + sum(pipe.flow for pipe in flowing),
+                            flowing | {c_pipe},
+                        )
+
+                if e_pipe in relevant_pipes and e_pipe not in flowing:
+                    for c_next_s in c_pipe.tunnels:
+
+                        c_next = pipes[c_next_s]
+
+                        update_with_better(
+                            node_at_times[time + 1][c_next, e_pipe],
+                            flow + sum(pipe.flow for pipe in flowing),
+                            flowing | {e_pipe},
+                        )
+
+                if (
+                    e_pipe in relevant_pipes
+                    and c_pipe in relevant_pipes
+                    and e_pipe not in flowing
+                    and c_pipe not in flowing
+                ):
+                    update_with_better(
+                        node_at_times[time + 1][c_pipe, e_pipe],
+                        flow + sum(pipe.flow for pipe in flowing),
+                        flowing | {c_pipe, e_pipe},
+                    )
+
+                update_with_better(
+                    node_at_times[max_time][c_pipe, e_pipe],
+                    flow + sum(pipe.flow for pipe in flowing) * (max_time - time),
+                    flowing,
+                )
+
+                d3 += time_p.time() - t1
+
+    return max(
+        flow
+        for nodes_of_pipe in node_at_times[max_time].values()
+        for flow in nodes_of_pipe.values()
+    )
 
 
 # === MAIN ===
@@ -159,4 +267,4 @@ relevant_pipes = frozenset(pipe for pipe in pipes.values() if pipe.flow > 0)
 print(part_1(pipes["AA"], 30, distances, relevant_pipes))
 
 # 1707, 2223
-print(part_2(pipes["AA"], 26, distances, relevant_pipes))
+print(part_2(pipes["AA"], 26, pipes, relevant_pipes, distances))