Advertisement
viligen

chain_lightning_modified_Prim

Aug 11th, 2022
794
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.90 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. def calc_damage(jumps, damage):
  15.     for _ in range(jumps):
  16.         damage //= 2
  17.     return damage
  18.  
  19.  
  20. def prim(node, damage, damage_by_node, graph):
  21.     damage_by_node[node] += damage
  22.  
  23.     tree = {node}
  24.     jumps = [0] * len(graph)
  25.     pq = PriorityQueue()
  26.     for edge in graph[node]:
  27.         pq.put(edge)
  28.  
  29.     while not pq.empty():
  30.         min_edge = pq.get()
  31.         none_tree_node = -1
  32.         tree_node = -1
  33.  
  34.         if min_edge.first in tree and min_edge.second not in tree:
  35.             none_tree_node = min_edge.second
  36.             tree_node = min_edge.first
  37.         elif min_edge.first not in tree and min_edge.second in tree:
  38.             none_tree_node = min_edge.first
  39.             tree_node = min_edge.second
  40.  
  41.         if none_tree_node == -1:
  42.             continue
  43.  
  44.         tree.add(none_tree_node)
  45.  
  46.         for edge in graph[none_tree_node]:
  47.             pq.put(edge)
  48.  
  49.         jumps[none_tree_node] = jumps[tree_node] + 1
  50.         damage_by_node[none_tree_node] += calc_damage(jumps[none_tree_node], damage)
  51.  
  52.  
  53. nodes = int(input())
  54. edges = int(input())
  55. lightnings = int(input())
  56.  
  57. graph = {node: [] for node in range(nodes)}
  58.  
  59. for _ in range(edges):
  60.     first, second, weight = [int(x) for x in input().split(' ')]
  61.     if first not in graph:
  62.         graph[first] = []
  63.     if second not in graph:
  64.         graph[second] = []
  65.     edge = Edge(first, second, weight)
  66.     graph[first].append(edge)
  67.     graph[second].append(edge)
  68.  
  69. damage_by_node = [0] * nodes
  70. for _ in range(lightnings):
  71.     hit_node, damage = [int(x) for x in input().split(' ')]
  72.     prim(hit_node, damage, damage_by_node, graph)
  73.  
  74. print(max(damage_by_node))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement