代码随想录算法训练营第五十七天｜prim算法、kruskal算法

代码随想录prim算法：

def prim(grid, v):
    # 初始化最小距离数组，设置为无穷大
    min_dist = [float('inf')] * (v + 1)
    # 初始化是否在树中的标记数组
    is_in_tree = [False] * (v + 1)

    # 循环 v-1 次，构建最小生成树
    for _ in range(v - 1):
        # 选择距离生成树最近节点
        cur, min_val = -1, float('inf')
        for j in range(1, v + 1):
            if not is_in_tree[j] and min_dist[j] < min_val:
                min_val = min_dist[j]
                cur = j
        # 将最近节点加入生成树
        is_in_tree[cur] = True

        # 更新非生成树节点到生成树的距离
        for j in range(1, v + 1):
            if not is_in_tree[j] and grid[cur][j] < min_dist[j]:
                min_dist[j] = grid[cur][j]

    # 计算最小生成树的总权重
    result = sum(min_dist[1:-1])  # 不计第一个和最后一个顶点
    return result

def main():
    v, e = map(int, input().split())
    # 初始化邻接矩阵，设置为默认最大值
    grid = [[10001 for _ in range(v + 1)] for _ in range(v + 1)]
    while e:
        x, y, k = map(int, input().split())
        # 因为是双向图，所以两个方向都要填上
        grid[x][y] = k
        grid[y][x] = k
        e -= 1

    # 调用 Prim 算法求最小生成树的总权重
    result = prim(grid, v)
    print(result)

if __name__ == "__main__":
    main()

kruskal算法

代码随想录

class UnionFind:
    def __init__(self, n):
        self.father = [-1] * n

    def find(self, u):
        if u != self.father[u]:
            self.father[u] = self.find(self.father[u])
        return self.father[u]

    def union(self, u, v):
        x = self.find(u)
        y = self.find(v)
        if x == y:
            return
        self.father[y] = x

def kruskal(edges, n):
    uf = UnionFind(n)
    result_val = 0
    result = []
    edges.sort(key=lambda x: x.val)  # 按照边的权重排序

    for edge in edges:
        l, r, val = edge.l, edge.r, edge.val
        if uf.find(l) != uf.find(r):  # 如果两个端点不是在同一个集合
            result.append(edge)  # 将边添加到结果中
            result_val += val  # 更新结果的总权重
            uf.union(l, r)  # 合并两个顶点的集合

    return result, result_val


def main():
    v, e = map(int, input().split())
    edges = []
    while e:
        v1, v2, val = map(int, input().split())
        edges.append(Edge(v1, v2, val))
        e -= 1

    result, result_val = kruskal(edges, v)

    # 打印最小生成树的边
    for edge in result:
        print(f"{edge.l} - {edge.r} : {edge.val}")

    print("Total minimum cost:", result_val)


class Edge:
    def __init__(self, l, r, val):
        self.l = l
        self.r = r
        self.val = val


if __name__ == "__main__":
    main()