Advertisement
viligen

cable_network_Prim

Aug 9th, 2022
790
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.38 KB | None | 0 0
  1. from queue import PriorityQueue
  2.  
  3.  
  4. class Edge:
  5.     def __init__(self, first, second, weight):
  6.         self.first = first
  7.         self.second = second
  8.         self.weight = weight
  9.  
  10.     def __gt__(self, other):
  11.         return self.weight > other.weight
  12.  
  13.  
  14. budget = int(input())
  15. nodes = int(input())
  16. edges = int(input())
  17.  
  18. graph = [[] for _ in range(nodes)]
  19. tree = set()
  20.  
  21. for _ in range(edges):
  22.     edge_data = input().split()
  23.     first, second, weight = int(edge_data[0]), int(edge_data[1]), int(edge_data[2])
  24.     graph[first].append(Edge(first, second, weight))
  25.     graph[second].append(Edge(first, second, weight))
  26.  
  27.     if len(edge_data) == 4:
  28.         tree.add(first)
  29.         tree.add(second)
  30.  
  31. pq = PriorityQueue()
  32. for node in tree:
  33.     for edge in graph[node]:
  34.         pq.put(edge)
  35.  
  36. budget_used = 0
  37. while not pq.empty():
  38.     min_edge = pq.get()
  39.     none_tree_node = None
  40.  
  41.     if min_edge.first in tree and min_edge.second not in tree:
  42.         none_tree_node = min_edge.second
  43.     elif min_edge.first not in tree and min_edge.second in tree:
  44.         none_tree_node = min_edge.first
  45.  
  46.     if none_tree_node is None:
  47.         continue
  48.  
  49.     if budget_used + min_edge.weight > budget:
  50.         break
  51.  
  52.     budget_used += min_edge.weight
  53.     tree.add(none_tree_node)
  54.     for edge in graph[none_tree_node]:
  55.         pq.put(edge)
  56.  
  57. print(f'Budget used: {budget_used}')
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement