From a1e634858f8936f4a773f927f009fc57e44ca81c Mon Sep 17 00:00:00 2001
From: I A <ia767@nyu.edu>
Date: Sat, 17 Dec 2022 09:18:00 -0500
Subject: [PATCH] [2022][day 16] add part 2 solution

---
 2022/AOC16.py | 83 +++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 80 insertions(+), 3 deletions(-)

diff --git a/2022/AOC16.py b/2022/AOC16.py
index 72a4e78..0b70f56 100644
--- a/2022/AOC16.py
+++ b/2022/AOC16.py
@@ -101,18 +101,95 @@ def find_path(t: int, valve: str, state: Tuple[str, int], seen) -> int:
 
         return max(wait, max_new_valve)
 
-    return find_path(27, "AA", (), {})
+    return find_path(30, "AA", (), {})
+
+
+def part2(valves, network, max_t=26):
+    def find_path(t: int, valve_a: str, valve_b: str, state, seen) -> int:
+        current_pressure = sum(
+            (opened_time - t) * valves[v] for v, opened_time in state
+        )
+        if seen.get((t, valve_a, valve_b), -1) >= current_pressure:
+            return 0
+        seen[t, valve_a, valve_b] = current_pressure
+
+        if t <= 0 or len(state) == len(valves):
+            return current_pressure
+
+        options = [0]
+        opened_valves = [v for v, _ in state]
+
+        # consider all options
+        for new_valve_a in network[valve_a] + [valve_a]:
+            # player A opens current valve
+            if (
+                valve_a == new_valve_a
+                and valve_a not in opened_valves
+                and valves[valve_a] > 0
+            ):
+                for new_valve_b in network[valve_b] + [valve_b]:
+                    # update state
+                    new_state = [(v, tt) for v, tt in state]
+                    new_state.append((new_valve_a, t))
+
+                    # A - player E moves to new valve
+                    future_value = find_path(
+                        t - 1, new_valve_a, new_valve_b, new_state, seen
+                    )
+                    options.append(future_value)
+
+                    # B - player E opens current valve
+                    if (
+                        valve_b == new_valve_b
+                        and new_valve_b != new_valve_a
+                        and new_valve_b not in opened_valves
+                        and valves[new_valve_b] > 0
+                    ):
+                        new_state.append((new_valve_b, t))
+                        future_value = find_path(
+                            t - 1, new_valve_a, new_valve_b, new_state, seen
+                        )
+                        options.append(future_value)
+                        new_state.pop()
+            # player A moves to new valve
+            else:
+                for new_valve_b in network[valve_b] + [valve_b]:
+                    # A - player E moves to new valve
+                    future_value = find_path(
+                        t - 1, new_valve_a, new_valve_b, state, seen
+                    )
+                    options.append(future_value)
+
+                    # B - player E opens current valve
+                    if (
+                        valve_b == new_valve_b
+                        and new_valve_b != new_valve_a
+                        and new_valve_b not in opened_valves
+                        and valves[new_valve_b] > 0
+                    ):
+                        # update state
+                        new_state = [(v, tt) for v, tt in state]
+                        new_state.append((new_valve_b, t))
+                        future_value = find_path(
+                            t - 1, new_valve_a, new_valve_b, new_state, seen
+                        )
+                        options.append(future_value)
+                        new_state.pop()
+
+        return max(options)
+
+    return find_path(max_t - 1, "AA", "AA", (), {})
 
 
 @timer
 def main(filename: str) -> Solution:
     valves, network = process_input(filename)
     part_1_solution = part1(valves, network)
-    part_2_solution = 0
+    part_2_solution = part2(valves, network)
     return part_1_solution, part_2_solution
 
 
 if __name__ == "__main__":
     part_1_solution, part_2_solution = main("aoc16.txt")
     print(f"PART 1: {part_1_solution}")  # 1789
-    print(f"PART 2: {part_2_solution}")  #
+    print(f"PART 2: {part_2_solution}")  # 2496