Python中使用Kruskal算法求最小生成树

# 定义边
class Edge(object):
    def __init__(self, start, end, weight):
        self.start = start
        self.end = end
        self.weight = weight

# 定义图
class Graph(object):
    def __init__(self, vertices):
        self.V = vertices
        self.edges = []

    def add_edge(self, start, end, weight):
        self.edges.append(Edge(start, end, weight))

# 定义并查集
class UnionFind(object):
    def __init__(self, vertices):
        self.parent = {}
        for v in vertices:
            self.parent[v] = v

    def find(self, v):
        if self.parent[v] == v:
            return v
        self.parent[v] = self.find(self.parent[v])
        return self.parent[v]

    def union(self, v1, v2):
        p1 = self.find(v1)
        p2 = self.find(v2)
        self.parent[p1] = p2

# Kruskal算法求最小生成树
def kruskal(g):
    uf = UnionFind(g.V)
    g.edges = sorted(g.edges, key=lambda e: e.weight)
    mst_weight = 0
    mst_edges = []
    for e in g.edges:
        start_root = uf.find(e.start)
        end_root = uf.find(e.end)
        if start_root == end_root:
            continue
        uf.union(e.start, e.end)
        mst_weight += e.weight
        mst_edges.append(e)
    return mst_weight, mst_edges

# 示例
g = Graph(["p", "i", "d", "a", "n", "c", "o", "d", "e", "m"])
g.add_edge("p", "i", 1)
g.add_edge("p", "d", 3)
g.add_edge("i", "d", 2)
g.add_edge("i", "a", 5)
g.add_edge("d", "a", 4)
g.add_edge("a", "n", 8)
g.add_edge("a", "c", 7)
g.add_edge("a", "m", 3)
g.add_edge("n", "c", 6)
g.add_edge("c", "o", 2)
g.add_edge("c", "m", 9)
print(kruskal(g))