# Untitled

unknown
python
2 years ago
3.3 kB
1
Indexable
Never
```from advent_of_code.lib import parse as aoc_parse
from advent_of_code.lib import search as aoc_search
from collections import defaultdict, deque
from itertools import combinations
import re

def get_input(file):
# graph will be {valve : list of connected valves}
graph = defaultdict(list)
# {valve : flow rate}
flow_rates = {}
for line in aoc_parse.as_lines(aoc.read_input(2022, 16, file)):
res = re.search(r"Valve (..) has flow rate=(\d+); tunnels? leads? to valves? (.+)", line)
name = res.group(1)
flow_rates[name] = int(res.group(2))
graph[name] = res.group(3).split(", ")
return graph, flow_rates

def compute_solutions(dists, flow_rates, time):
# solutions is a dictionary {set of open valves: greatest pressure achievable}
solutions = defaultdict(int)
# node = position, time left, open valves, pressure
start = "AA", time, frozenset(), 0
todo = deque([start])
# {(position, open valves) : pressure}
visited = defaultdict(int)
# why is BFS faster than DFS?
while todo:
valve, time, open, pressure = todo.popleft()

# store even if not finished because in part2 will need non intersecting solutions
# which wouldn't be considered otherwise
solutions[open] = max(solutions[open], pressure)

# at least 1 to move and 1 to open new valve
if time <= 2:
continue

# move to every possible other valve and open it
for new_valve in dists[valve]:
# unless has flow 0 or already open
if new_valve in open or flow_rates[new_valve] == 0:
continue
new_open = open | {new_valve}
new_time = time - dists[valve][new_valve] - 1
new_pressure = pressure + new_time * flow_rates[new_valve]
# discard if no time left or state already visited with higher pressure
# where state is current valve and set of open valves
if new_time > 0 and visited[(new := (new_valve, new_open))] < new_pressure:
visited[new] = new_pressure
todo.append((new_valve, new_time, new_open, new_pressure))
return solutions

@aoc.pretty_solution(1)
def part1(graph, flow_rates):
dists = aoc_search.floyd_warshall(graph)
solutions = compute_solutions(dists, flow_rates, 30)
return max(solutions.values())

@aoc.pretty_solution(2)
def part2(graph, flow_rates):
dists = aoc_search.floyd_warshall(graph)
# basically the two agents are independent because they open different valves.
# Compute solution for single agent with 26 minutes then just take the max
# possible sum of the two pressures achieved independently by the agents
solutions = compute_solutions(dists, flow_rates, 26)
return max(p1 + p2 for (s1, p1), (s2, p2) in combinations(solutions.items(), 2) if not s1 & s2)

def main():
data = get_input("input.txt")
part1(*data)
part2(*data)

def test():
example = get_input("example.txt")
assert part1(*example) == 1651
assert part2(*example) == 1707

data = get_input("input.txt")
assert part1(*data) == 2114
assert part2(*data) == 2666

print("Test OK")

if __name__ == "__main__":
test()
```