Advertisement
cyberjab

Untitled

May 22nd, 2023
86
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.83 KB | None | 0 0
  1. import networkx as nx
  2. import matplotlib.pyplot as plt
  3.  
  4.  
  5. def kruskal(G, mx=False):
  6.     sp = []
  7.     n = len(G.nodes.items())
  8.     for (u, v, wt) in G.edges.data('weight'):
  9.         sp.append((u, v, wt))
  10.     if not mx:
  11.         sp.sort(key=lambda x: x[2])
  12.         comp = [i for i in range(n)]
  13.         ans = 0
  14.         spp = {}
  15.         x = []
  16.         for start, end, weight in sp:
  17.             start -= 1
  18.             end -= 1
  19.             if comp[start] != comp[end]:
  20.                 ans += weight
  21.                 a = comp[start]
  22.                 b = comp[end]
  23.                 for i in range(n):
  24.                     if comp[i] == b:
  25.                         comp[i] = a
  26.                 spp[(start + 1, end + 1)] = weight
  27.                 x.append((start + 1, end + 1))
  28.         return spp, ans, x
  29.     else:
  30.         sp.sort(key=lambda x: -x[2])
  31.         comp = [i for i in range(n)]
  32.         ans = 0
  33.         spp = {}
  34.         x = []
  35.         for start, end, weight in sp:
  36.             start -= 1
  37.             end -= 1
  38.             if comp[start] != comp[end]:
  39.                 ans += weight
  40.                 a = comp[start]
  41.                 b = comp[end]
  42.                 for i in range(n):
  43.                     if comp[i] == b:
  44.                         comp[i] = a
  45.                 spp[(start + 1, end + 1)] = weight
  46.                 x.append((start + 1, end + 1))
  47.         return spp, ans, x
  48.  
  49.  
  50. G = nx.Graph()
  51.  
  52. with open("input_graph") as f:
  53.     n, m = [int(x) for x in f.readline().strip().split()]
  54.     for i in range(n):
  55.         posx, posy = [int(x) for x in f.readline().strip().split()]
  56.         G.add_node(i + 1, pos=(posx, posy))
  57.     for i in range(m):
  58.         first_point, second_point, w = [int(x) for x in f.readline().strip().split()]
  59.         G.add_edge(first_point, second_point, weight=w)
  60.  
  61. # G.add_node(1, pos=(0, 4))
  62. # G.add_node(2, pos=(2, 3))
  63. # G.add_node(3, pos=(4, 0))
  64. # G.add_node(4, pos=(2, -3))
  65. # G.add_node(5, pos=(0, -4))
  66. # G.add_node(6, pos=(-3, -2))
  67. # G.add_node(7, pos=(-4, 0))
  68. # G.add_node(8, pos=(-3, 2))
  69. #
  70. # G.add_edge(1, 2, weight=2)
  71. # G.add_edge(1, 8, weight=1)
  72. #
  73. # G.add_edge(2, 4, weight=1)
  74. # G.add_edge(2, 6, weight=3)
  75. # G.add_edge(2, 8, weight=3)
  76. #
  77. # G.add_edge(3, 4, weight=2)
  78. # G.add_edge(3, 7, weight=6)
  79. #
  80. # G.add_edge(4, 5, weight=2)
  81. # G.add_edge(4, 8, weight=3)
  82. #
  83. # G.add_edge(5, 6, weight=1)
  84.  
  85. sp, ans, spp = kruskal(G)
  86.  
  87. pos = nx.get_node_attributes(G, 'pos')
  88. labels = nx.get_edge_attributes(G, 'weight')
  89.  
  90. nx.draw(G, pos, with_labels=True, font_weight='bold', font_color="white", font_size=14)
  91. nx.draw_networkx_nodes(G, pos, node_size=400, node_color="black")
  92.  
  93. nx.draw_networkx_edges(G, pos, edge_color="blue", edgelist=spp)
  94. nx.draw_networkx_edge_labels(G, pos, edge_labels=labels, font_size=16)
  95. nx.draw_networkx_edge_labels(G, pos, edge_labels=sp, font_color="blue", font_size=16)
  96.  
  97. plt.show()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement