Greedy Algorithms - Prim

from pyvis.network import Network
from heapq import heappush
from heapq import heappop


def prim(edges):
    connected_edges = dict()
    for u, v, weight in edges:
        if u not in connected_edges:
            connected_edges[u] = set()
        if v not in connected_edges:
            connected_edges[v] = set()
        connected_edges[u].add((u, v, weight))
        connected_edges[v].add((v, u, weight))

    next_edges = list() # heap of next edges
    visited = set() # visited nodes
    heappush(next_edges, (0, None, edges[0][0]))
    mst_edges = list()

    #edges[0][0]
    while next_edges:
        next_edge = heappop(next_edges)
        # Αποφασίζω ότι στην αρχή θα είναι το βάρος και μετά οι 2 κορυφές
        next_edge_weight = next_edge[0]
        next_edge_u = next_edge[1]
        next_edge_v = next_edge[2]

        if next_edge_v in visited:
            continue
        visited.add(next_edge_v)
        if next_edge_u is not None:
            mst_edges.append((next_edge_u, next_edge_v, next_edge_weight))
        for u, v, weight in connected_edges[next_edge_v]:
            if v not in visited:
                # Εδώ, γράφω ότι πρώτα στο tuple θα μπει το βάρος, για να γίνει
                # μετά σωστά το pus, pop βάσει των 1ων στοιχείων των tuples =  των βαρών
                heappush(next_edges, (weight, u, v))
    return mst_edges



def visualize(edges):
    G = Network()
    for u, v, weight in edges:
        G.add_node(u)
        G.add_node(v)
        G.add_edge(u, v, label=str(weight))
    G.show("network.html")

# MAIN FUNCTION
edges = [('A', 'B', 2), ('A', 'C', 3), ('A', 'D', 3), ('B', 'C', 4),
      ('B', 'E', 3), ('C', 'D', 5), ('C', 'E', 1), ('C', 'F', 6),
      ('D', 'F', 7), ('E', 'F', 8), ('F', 'G', 9)]
mst_edges = prim(edges)
visualize(mst_edges)