Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import heapq
- def dijkstra(valves, source):
- pq = [(0, float('inf'), source)] # (total flow rate, distance, valve)
- dist = {source: 0} # distance from source
- flow = {source: float('inf')} # total flow rate from source
- prev = {} # previous valve in the path
- while pq:
- flow_rate, distance, valve = heapq.heappop(pq)
- if distance > 30: # stop searching after 30 minutes
- break
- if flow_rate < flow[valve]:
- continue # this valve has already been visited with higher flow rate
- for neighbor, flow_rate_neighbor in valves[valve]:
- new_flow_rate = min(flow[valve], flow_rate_neighbor)
- new_distance = distance + 1
- if neighbor not in dist or new_flow_rate > flow[neighbor]:
- flow[neighbor] = new_flow_rate
- dist[neighbor] = new_distance
- prev[neighbor] = valve
- heapq.heappush(pq, (new_flow_rate, new_distance, neighbor))
- return prev, flow
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement