Kruskal method

kruskal.py

def union_find(u, arr):
    if arr[u] == u:
        return u
    return union_find(arr[u], arr)

def union(u, v, arr):    
    arr[union_find(v, arr)] = union_find(u, arr)
    
def kruskal(nodes, edges):
    mst = []
    edges = sorted(edges, key=lambda e: e[2])        
    arr = [i for i in range(0, len(nodes))]    
    base = ord('A') # for indexing the vertice
    for e in edges:
        u, v, w = e
        u = int(ord(u) - base)
        v = int(ord(v) - base)
        if union_find(u, arr) != union_find(v, arr):
            mst.append(e)
            union(u, v, arr)         
    return mst


#test
nodes = list("ABCDEFG")
edges = [ ("A", "B", 7), ("A", "D", 5),
          ("B", "C", 8), ("B", "D", 9), ("B", "E", 7),
      ("C", "E", 5),
      ("D", "E", 15), ("D", "F", 6),
      ("E", "F", 8), ("E", "G", 9),
      ("F", "G", 11)]

print("kruskal:", kruskal(nodes, edges))