Advertisement
cyberjab

Untitled

May 24th, 2023
126
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.91 KB | None | 0 0
  1. import networkx as nx
  2. import matplotlib.pyplot as plt
  3.  
  4.  
  5. def kruskal(G, make, mx=False):
  6.     if make:
  7.         sp = []
  8.         n = len(G.nodes.items())
  9.         for (u, v, wt) in G.edges.data('weight'):
  10.             sp.append((u, v, wt))
  11.         if not mx:
  12.             sp.sort(key=lambda x: x[2])
  13.             comp = [i for i in range(n)]
  14.             ans = 0
  15.             spp = {}
  16.             x = []
  17.             for start, end, weight in sp:
  18.                 start -= 1
  19.                 end -= 1
  20.                 if comp[start] != comp[end]:
  21.                     ans += weight
  22.                     a = comp[start]
  23.                     b = comp[end]
  24.                     for i in range(n):
  25.                         if comp[i] == b:
  26.                             comp[i] = a
  27.                     spp[(start + 1, end + 1)] = weight
  28.                     x.append((start + 1, end + 1))
  29.             return spp, ans, x
  30.         else:
  31.             sp.sort(key=lambda x: -x[2])
  32.             comp = [i for i in range(n)]
  33.             ans = 0
  34.             spp = {}
  35.             x = []
  36.             for start, end, weight in sp:
  37.                 start -= 1
  38.                 end -= 1
  39.                 if comp[start] != comp[end]:
  40.                     ans += weight
  41.                     a = comp[start]
  42.                     b = comp[end]
  43.                     for i in range(n):
  44.                         if comp[i] == b:
  45.                             comp[i] = a
  46.                     spp[(start + 1, end + 1)] = weight
  47.                     x.append((start + 1, end + 1))
  48.             return spp, ans, x
  49.     return (0, 0, 0)
  50.  
  51.  
  52. G = nx.Graph()
  53. make_graph = True
  54. with open("inputkr.txt") as f:
  55.     sp = f.readlines()
  56.     n, m = [int(x) for x in sp[0].strip().split()]
  57.     check = len(sp)
  58.     n_cnt = 0
  59.     m_cnt = 0
  60.     for i in range(1, check):
  61.         s = [int(x) for x in sp[i].strip().split()]
  62.         if len(s) == 2:
  63.             n_cnt += 1
  64.         if len(s) == 3:
  65.             m_cnt += 1
  66.     if n_cnt < n:
  67.         print("Введены не все вершины")
  68.         make_graph = False
  69.     if n_cnt > n:
  70.         print("Введено больше вершин чем указано")
  71.         make_graph = False
  72.     if m_cnt < m:
  73.         print("Введены не все ребра")
  74.         make_graph = False
  75.     if m_cnt > m:
  76.         print("Введено больше ребер чем указано")
  77.         make_graph = False
  78.     if make_graph:
  79.         for i in range(1, n + 1):
  80.             posx, posy = [int(x) for x in sp[i].strip().split()]
  81.             G.add_node(i, pos=(posx, posy))
  82.         for i in range(n + 1, check):
  83.             first_point, second_point, w = [int(x) for x in sp[i].strip().split()]
  84.             G.add_edge(first_point, second_point, weight=w)
  85.  
  86. # G.add_node(1, pos=(0, 4))
  87. # G.add_node(2, pos=(2, 3))
  88. # G.add_node(3, pos=(4, 0))
  89. # G.add_node(4, pos=(2, -3))
  90. # G.add_node(5, pos=(0, -4))
  91. # G.add_node(6, pos=(-3, -2))
  92. # G.add_node(7, pos=(-4, 0))
  93. # G.add_node(8, pos=(-3, 2))
  94. #
  95. # G.add_edge(1, 2, weight=2)
  96. # G.add_edge(1, 8, weight=1)
  97. #
  98. # G.add_edge(2, 4, weight=1)
  99. # G.add_edge(2, 6, weight=3)
  100. # G.add_edge(2, 8, weight=3)
  101. #
  102. # G.add_edge(3, 4, weight=2)
  103. # G.add_edge(3, 7, weight=6)
  104. #
  105. # G.add_edge(4, 5, weight=2)
  106. # G.add_edge(4, 8, weight=3)
  107. #
  108. # G.add_edge(5, 6, weight=1)
  109.  
  110. sp, ans, spp = kruskal(G, make_graph)
  111.  
  112. if make_graph:
  113.     pos = nx.get_node_attributes(G, 'pos')
  114.     labels = nx.get_edge_attributes(G, 'weight')
  115.  
  116.     nx.draw(G, pos, with_labels=True, font_weight='bold', font_color="white", font_size=14)
  117.     nx.draw_networkx_nodes(G, pos, node_size=400, node_color="black")
  118.  
  119.     nx.draw_networkx_edges(G, pos, edge_color="blue", edgelist=spp)
  120.     nx.draw_networkx_edge_labels(G, pos, edge_labels=labels, font_size=16)
  121.     nx.draw_networkx_edge_labels(G, pos, edge_labels=sp, font_color="blue", font_size=16)
  122.  
  123.     plt.show()
  124. else:
  125.     print("Граф задан некорректно")
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement