Coverage for src\aoc2022\days\day16.py: 100%
101 statements
« prev ^ index » next coverage.py v7.0.0, created at 2022-12-23 05:19 +0000
1import itertools
2from typing import Dict, Generator, List, Optional, Set, Tuple
4from src.aoc2022 import utils
7class Graph:
8 """
9 Data class representing a graph of nodes and edges
11 :param flow_rates: A dict of node name, flow rate
12 :type flow_rates: Dict[str, int]
14 :param cons: a dict of node name, node edges
15 :type cons: Dict[str, List[str]]
16 """
18 def __init__(self, flow_rates: Dict[str, int], cons: Dict[str, List[str]]):
19 self.flow_rates = flow_rates
20 self.valves = sorted(list(self.flow_rates.keys()))
21 self.connections = cons
22 self.distance_matrix = {}
23 for node in flow_rates.keys():
24 distances = self.compute_distance_matrix(node)
25 self.distance_matrix[node] = distances
27 def get_valves(self) -> List[str]:
28 """
29 Get a list of all valve names
31 :return: valve names
32 :rtype: List[str]
33 """
34 return self.valves
36 def get_connections(self, node_name: str) -> List[str]:
37 """
38 Get connections for a specified node
40 :param node_name: Name of node to retrieve connections
41 :type node_name: str
43 :return: the names of the connected nodes for the given node
44 :rtype: List[str]
45 """
46 return self.connections[node_name]
48 def get_flow_rate(self, node_name) -> int:
49 """
50 Get flow rate for a specified node
52 :param node_name: Name of node to retrieve flow rate
53 :type node_name: str
55 :return: Flow rate of given node
56 :rtype: int
57 """
58 return self.flow_rates[node_name]
60 def get_flow_rates(self) -> Dict[str, int]:
61 """
62 Get flow rates for all nodes
64 :return: Dict of node name, flow rate
65 :rtype: Dict[str, int]
66 """
67 return self.flow_rates
69 def get_minimal_distance(self, start: str, end: str) -> int:
70 """
71 Get shortest distance between two valves
73 :param start: Name of starting node
74 :type start: str
76 :param node_name: Name of end node
77 :type node_name: str
79 :return: shortest distance between both nodes
80 :rtype: int
81 """
82 return self.distance_matrix[start][end]
84 def compute_distance_matrix(self, start: str) -> Dict[str, int]:
85 """
86 Compute shortest distance from starting node to all other nodes
88 :param start: Name of starting node
89 :type start: str
91 :return: shortest distance from starting node to all other nodes
92 :rtype: Dict[str, int]
93 """
94 distances = {start: 0}
96 stack = [start]
97 visited = set()
99 while stack:
100 node = stack.pop(0)
101 if node not in visited:
102 visited.add(node)
103 for neighbour in self.get_connections(node):
104 stack.append(neighbour)
105 increment_dist = distances[node] + 1
106 default_dist = distances.get(neighbour, increment_dist)
107 distances[neighbour] = min(default_dist, increment_dist)
109 return distances
111 def time_aware_permutations(
112 self, position: str, remaining_valves: Set[str], remaining_time: int
113 ) -> Generator[List[Tuple[str, int]], None, None]:
114 """
115 Generate all possible permutations of valve openings that respect
116 the time limits
118 :param position: Name of current node
119 :type position: str
121 :param remaining_valves: valves that have not yet been opened
122 :type remaining_valves: Set[str]
124 :param remaining_time: Amount of time remaining before eruption
125 :type remaining_time: int
127 :return: shortest distance from starting node to all other nodes
128 :rtype: Dict[str, int]
129 """
130 if (remaining_time <= 0) or (len(remaining_valves) == 0):
131 yield [(position, max(remaining_time, 0))]
132 else:
133 for new_pos in remaining_valves:
134 valves = remaining_valves.copy()
135 valves.remove(new_pos)
137 dist = self.get_minimal_distance(position, new_pos)
139 updated_time = remaining_time - dist - 1
141 for subsearch in self.time_aware_permutations(
142 new_pos, valves, updated_time
143 ):
144 yield [(position, remaining_time)] + subsearch
147def parse_data_to_graph(raw_data: str) -> Graph:
148 """
149 Parse raw data into a Graph of valves and connections
151 :param raw_data: Raw data to parse into Graph
152 :type raw_data: str
154 :return: the parsed Graph data
155 :rtype: Graph
156 """
158 connections = {}
159 flow_rates = {}
161 for line in raw_data.split("\n"):
162 flow_rate_str, connection_str = line.split("; ")
164 valve_name_str, flow_rate_str = flow_rate_str.split("rate=")
165 flow_rate = int(flow_rate_str)
166 valve_name = valve_name_str.split(" ")[1]
168 connection_str = connection_str.replace("valves", "valve")
169 valve_connections = connection_str.split("valve ")[-1].split(", ")
171 flow_rates[valve_name] = flow_rate
172 connections[valve_name] = valve_connections
174 return Graph(flow_rates, connections)
177def get_optimal_sequence(
178 graph: Graph,
179 start: str,
180 total_time: int,
181 non_zero_valves: Optional[List[str]] = None,
182) -> Tuple[List[Tuple[str, int]], int]:
183 """
184 Compute the optimal valve opening sequence
186 :param graph: Graph of valves to open
187 :type graph: Graph
189 :param start: Name of starting node
190 :type start: str
192 :param total_time: total time for valve sequence
193 :type total_time: int
195 :param non_zero_valves: Optional list of valves to specify for search
196 :type non_zero_valves: Optional[List[str]]
198 :return: Optimal sequence and score for opening valves
199 :rtype: Tuple[List[Tuple[str, int]], int]
200 """
201 if non_zero_valves is None:
202 valves = graph.get_flow_rates()
203 non_zero_valves = [k for k, v in valves.items() if v > 0]
205 max_ = 0
206 max_path = []
208 search_valves = set(non_zero_valves)
209 paths = graph.time_aware_permutations(start, search_valves, total_time)
211 for path in paths:
213 score = sum([v * graph.get_flow_rate(k) for k, v in path])
215 if score > max_:
216 max_ = score
217 max_path = path
219 return max_path, max_
222def get_optimal_dual_sequence(
223 graph: Graph, start: str, total_time: int
224) -> Tuple[List[Tuple[str, int]], List[Tuple[str, int]], int]:
225 """
226 Compute the optimal valve opening sequence for two agents
228 :param graph: Graph of valves to open
229 :type graph: Graph
231 :param start: Name of starting node
232 :type start: str
234 :param total_time: total time for valve sequence
235 :type total_time: int
237 :return: Optimal paths for elephant and person, and associated score
238 :rtype: Tuple[List[Tuple[str, int]], List[Tuple[str, int]], int]
239 """
240 real_valves = [k for k, v in graph.get_flow_rates().items() if v > 0]
241 n = len(real_valves)
243 optimal_cache: Dict[frozenset, Tuple] = {}
245 max_ = 0
246 path_type = List[Tuple[str, int]]
247 max_paths: Tuple[path_type, path_type] = ([], [])
249 partition_flag_combinations = itertools.product([True, False], repeat=n)
250 for idx, bool_flags in enumerate(partition_flag_combinations):
252 elephant_valves = [real_valves[i] for i in range(n) if bool_flags[i]]
253 person_valves = [real_valves[i] for i in range(n) if not bool_flags[i]]
255 elephant_cache_valves = frozenset(elephant_valves)
256 person_cache_valves = frozenset(person_valves)
258 e_score = 0
259 p_score = 0
261 if elephant_cache_valves in optimal_cache:
262 e_path, e_score = optimal_cache[elephant_cache_valves]
263 else:
264 e_path, e_score = get_optimal_sequence(
265 graph, start, total_time, non_zero_valves=elephant_valves
266 )
267 optimal_cache[elephant_cache_valves] = (e_path, e_score)
269 if person_cache_valves in optimal_cache:
270 p_path, p_score = optimal_cache[person_cache_valves]
271 else:
272 p_path, p_score = get_optimal_sequence(
273 graph, start, total_time, non_zero_valves=person_valves
274 )
275 optimal_cache[person_cache_valves] = (p_path, p_score)
277 combined_score = p_score + e_score
279 if combined_score > max_:
280 max_ = combined_score
281 max_paths = (e_path, p_path)
283 return *max_paths, max_
286if __name__ == "__main__": # pragma: no cover
288 raw_data = utils.get_raw_data("./src/aoc2022/data/day16.txt")
289 graph = parse_data_to_graph(raw_data)
290 start_node = "AA"
291 total_time = 30
293 path, score = get_optimal_sequence(graph, start_node, total_time)
295 print(f"The best path is {path} which will release {score} pressure")
297 total_time = 26
298 _, _, score = get_optimal_dual_sequence(graph, start_node, total_time)
300 print(f"The best path will release {score} pressure")