# 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 advent_of_code.lib import aoc 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()