分别用Prim算法与Kruskal算法求最小树的代码

Prim算法:

def prim(graph):
    """
    :param graph: 图，使用邻接矩阵表示
    :return: 最小生成树的边集合，使用列表存储
    """
    # 初始化
    num_vertices = len(graph)
    INF = float('inf')
    min_cost = [INF] * num_vertices
    visited = [False] * num_vertices
    parent = [-1] * num_vertices
    min_cost[0] = 0

    # 遍历所有节点
    for i in range(num_vertices):
        # 找到未访问过的最小权值的节点
        u = -1
        for j in range(num_vertices):
            if not visited[j] and (u == -1 or min_cost[j] < min_cost[u]):
                u = j
        visited[u] = True

        # 更新与节点u相邻的节点的权值
        for v in range(num_vertices):
            if graph[u][v] != INF and not visited[v] and graph[u][v] < min_cost[v]:
                min_cost[v] = graph[u][v]
                parent[v] = u

    # 构建最小生成树的边集合
    edges = []
    for i in range(1, num_vertices):
        edges.append((parent[i], i))

    return edges

Kruskal算法:

class DisjointSet:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

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

    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False

        if self.rank[root_x] < self.rank[root_y]:
            self.parent[root_x] = root_y
        elif self.rank[root_x] > self.rank[root_y]:
            self.parent[root_y] = root_x
        else:
            self.parent[root_y] = root_x
            self.rank[root_x] += 1

        return True


def kruskal(graph):
    """
    :param graph: 图，使用列表存储所有边
    :return: 最小生成树的边集合，使用列表存储
    """
    # 初始化
    num_vertices = len(set([v for e in graph for v in e]))
    disjoint_set = DisjointSet(num_vertices)
    edges = []
    min_cost = 0

    # 将图中所有边按照权值从小到大排序
    graph.sort(key=lambda x: x[2])

    # 遍历所有边
    for e in graph:
        u, v, weight = e
        if disjoint_set.union(u, v):
            edges.append((u, v))
            min_cost += weight

    return edges

注意：这里的图输入格式不同，Prim算法使用邻接矩阵表示，Kruskal算法使用边的列表表示。