Advertisement
makispaiktis

Greedy Algorithms - Prim

May 29th, 2020 (edited)
729
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.94 KB | None | 0 0
  1. from pyvis.network import Network
  2. from heapq import heappush
  3. from heapq import heappop
  4.  
  5.  
  6. def prim(edges):
  7.     connected_edges = dict()
  8.     for u, v, weight in edges:
  9.         if u not in connected_edges:
  10.             connected_edges[u] = set()
  11.         if v not in connected_edges:
  12.             connected_edges[v] = set()
  13.         connected_edges[u].add((u, v, weight))
  14.         connected_edges[v].add((v, u, weight))
  15.  
  16.     next_edges = list() # heap of next edges
  17.     visited = set() # visited nodes
  18.     heappush(next_edges, (0, None, edges[0][0]))
  19.     mst_edges = list()
  20.  
  21.     #edges[0][0]
  22.     while next_edges:
  23.         next_edge = heappop(next_edges)
  24.         # Αποφασίζω ότι στην αρχή θα είναι το βάρος και μετά οι 2 κορυφές
  25.         next_edge_weight = next_edge[0]
  26.         next_edge_u = next_edge[1]
  27.         next_edge_v = next_edge[2]
  28.  
  29.         if next_edge_v in visited:
  30.             continue
  31.         visited.add(next_edge_v)
  32.         if next_edge_u is not None:
  33.             mst_edges.append((next_edge_u, next_edge_v, next_edge_weight))
  34.         for u, v, weight in connected_edges[next_edge_v]:
  35.             if v not in visited:
  36.                 # Εδώ, γράφω ότι πρώτα στο tuple θα μπει το βάρος, για να γίνει
  37.                 # μετά σωστά το pus, pop βάσει των 1ων στοιχείων των tuples =  των βαρών
  38.                 heappush(next_edges, (weight, u, v))
  39.     return mst_edges
  40.  
  41.  
  42.  
  43. def visualize(edges):
  44.     G = Network()
  45.     for u, v, weight in edges:
  46.         G.add_node(u)
  47.         G.add_node(v)
  48.         G.add_edge(u, v, label=str(weight))
  49.     G.show("network.html")
  50.  
  51. # MAIN FUNCTION
  52. edges = [('A', 'B', 2), ('A', 'C', 3), ('A', 'D', 3), ('B', 'C', 4),
  53.       ('B', 'E', 3), ('C', 'D', 5), ('C', 'E', 1), ('C', 'F', 6),
  54.       ('D', 'F', 7), ('E', 'F', 8), ('F', 'G', 9)]
  55. mst_edges = prim(edges)
  56. visualize(mst_edges)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement