153 lines
4.5 KiB
Python
153 lines
4.5 KiB
Python
from collections import deque
|
|
from itertools import product
|
|
from tools.aoc import AOCDay
|
|
from typing import Any
|
|
from tqdm import tqdm
|
|
|
|
|
|
class Valve:
|
|
def __init__(self, name: str, flowrate: int) -> None:
|
|
self.name = name
|
|
self.flowrate = flowrate
|
|
self.tunnels = []
|
|
|
|
|
|
def get_openable_valve_tunnels(valve: Valve, open_valves: set, time_remaining: int) -> set:
|
|
tunnels = set()
|
|
queue = deque()
|
|
visited = set()
|
|
queue.append((0, valve))
|
|
while queue:
|
|
d, v = queue.popleft()
|
|
if v.name in visited:
|
|
continue
|
|
visited.add(v.name)
|
|
if v.name not in open_valves and v.flowrate > 0 and d + 2 <= time_remaining:
|
|
tunnels.add((d, v))
|
|
for x in v.tunnels:
|
|
queue.append((d + 1, x))
|
|
|
|
return tunnels
|
|
|
|
|
|
def get_max_flow(valve: Valve, open_valves: set, time_remaining: int = 30) -> int:
|
|
max_flow = 0
|
|
|
|
ov = get_openable_valve_tunnels(valve, open_valves, time_remaining)
|
|
if time_remaining <= 0 or not ov:
|
|
return 0
|
|
|
|
for d, v in ov:
|
|
this_open_valves = open_valves.copy()
|
|
if d + 2 > time_remaining:
|
|
continue
|
|
this_open_valves.add(v.name)
|
|
this_open_flow = v.flowrate * (time_remaining - d - 1)
|
|
this_flow = get_max_flow(v, this_open_valves, time_remaining - d - 1)
|
|
if this_flow + this_open_flow > max_flow:
|
|
max_flow = this_flow + this_open_flow
|
|
|
|
return max_flow
|
|
|
|
|
|
def get_max_flow_double(valve1: Valve, valve2: Valve, open_valves: set, DP: dict, time_remaining_1: int = 26, time_remaining_2: int = 26, depth: int = 0) -> int:
|
|
dp_key = valve1.name + valve2.name + "%02d" % time_remaining_1 + "%02d" % time_remaining_2 + "".join(list(sorted(open_valves)))
|
|
dp_key2 = valve2.name + valve1.name + "%02d" % time_remaining_1 + "%02d" % time_remaining_2 + "".join(list(sorted(open_valves)))
|
|
if dp_key in DP:
|
|
return DP[dp_key]
|
|
|
|
if time_remaining_1 <= 0 and time_remaining_2 <= 0:
|
|
return 0
|
|
|
|
ov1 = get_openable_valve_tunnels(valve1, open_valves, time_remaining_1)
|
|
ov2 = get_openable_valve_tunnels(valve2, open_valves, time_remaining_2)
|
|
|
|
if not ov1 and not ov2:
|
|
return 0
|
|
|
|
if not ov1:
|
|
ov1 = {(99, valve1)}
|
|
|
|
if not ov2:
|
|
ov2 = {(99, valve2)}
|
|
|
|
permut = product(ov1, ov2)
|
|
|
|
if depth == 0:
|
|
pbar = tqdm(total=210)
|
|
|
|
max_flow = 0
|
|
for v1, v2 in permut:
|
|
if v1[1].name == v2[1].name:
|
|
continue
|
|
if depth == 0:
|
|
pbar.update(1)
|
|
this_open_values = open_valves.copy()
|
|
d1, tv1 = v1
|
|
d2, tv2 = v2
|
|
this_open_flow = 0
|
|
if d1 + 2 <= time_remaining_1:
|
|
this_open_values.add(tv1.name)
|
|
this_open_flow += tv1.flowrate * (time_remaining_1 - d1 - 1)
|
|
else:
|
|
tv1 = valve1
|
|
if d2 + 2 <= time_remaining_2:
|
|
this_open_values.add(tv2.name)
|
|
this_open_flow += tv2.flowrate * (time_remaining_2 - d2 - 1)
|
|
else:
|
|
tv2 = valve2
|
|
|
|
this_flow_rate = get_max_flow_double(tv1, tv2, this_open_values, DP, time_remaining_1 - d1 - 1, time_remaining_2 - d2 - 1, depth + 1)
|
|
if this_flow_rate + this_open_flow > max_flow:
|
|
max_flow = this_flow_rate + this_open_flow
|
|
|
|
if depth == 0:
|
|
pbar.close()
|
|
DP[dp_key] = max_flow
|
|
DP[dp_key2] = max_flow
|
|
return max_flow
|
|
|
|
|
|
class Day(AOCDay):
|
|
inputs = [
|
|
[
|
|
(1651, "input16_test"),
|
|
(1850, "input16"),
|
|
],
|
|
[
|
|
(1707, "input16_test"),
|
|
(2306, "input16"),
|
|
]
|
|
]
|
|
|
|
def get_valve_graph(self) -> Valve:
|
|
valves = {}
|
|
for line in self.getInput():
|
|
p = line.split(" ")
|
|
valve_name = p[1]
|
|
flowrate = int(p[4][5:-1])
|
|
tunnels = "".join(p[9:]).split(",")
|
|
valves[valve_name] = Valve(valve_name, flowrate)
|
|
valves[valve_name].tunnels = tunnels
|
|
|
|
for valve in valves.values():
|
|
valve.tunnels = [valves[x] for x in valve.tunnels]
|
|
for valve in valves.values():
|
|
for v in valve.tunnels:
|
|
if valve not in v.tunnels:
|
|
v.tunnels.append(valve)
|
|
|
|
return valves["AA"]
|
|
|
|
def part1(self) -> Any:
|
|
return get_max_flow(self.get_valve_graph(), set())
|
|
|
|
def part2(self) -> Any:
|
|
root = self.get_valve_graph()
|
|
return get_max_flow_double(root, root, set(), {})
|
|
|
|
|
|
if __name__ == '__main__':
|
|
day = Day(2022, 16)
|
|
day.run(verbose=True)
|