python实现Kruskal && Prim算法

一、Kruskal算法

def find(forest, item):
    """返回未变化的双亲节点，如果变化了找到未变化的为止"""
    if forest[item] != item: # 双亲节点和当前节点不一致，变化了说明这条边已经添加MST
        forest[item] = find(forest, forest[item])  # 找到双亲节点未变化的点
    return forest[item]


def Kruskal(nodes, edges):
    '''基于不相交集实现Kruskal算法'''
    forest = {}  
    MST = []
    for item in nodes:  # 双亲节点初始化
        forest[item] = item

    edges = sorted(edges, key=lambda element: element[2])  # 排序
    num_sides = len(nodes) - 1  # 最小生成树的边数等于顶点数减一
    for e in edges:
        node1, node2, _ = e
        parent1 = find(forest, node1) # 找到node1的双亲节点
        parent2 = find(forest, node2)
        if parent1 != parent2:     # 双亲节点不一致
            MST.append(e)
            num_sides -= 1
            if num_sides == 0:     # 注意控制条件是最小生成树的边满足条件
                return MST
            else:                  # 改变该顶点的双亲节点
                forest[parent1] = parent2
   


def main():
    nodes = set(list('ABCDEFGHI'))  # 注意集合无序
    edges = [("A", "B", 4), ("A", "H", 8),
             ("B", "C", 8), ("B", "H", 11),
             ("C", "D", 7), ("C", "F", 4),
             ("C", "I", 2), ("D", "E", 9),
             ("D", "F", 14), ("E", "F", 10),
             ("F", "G", 2), ("G", "H", 1),
             ("G", "I", 6), ("H", "I", 7)]
    print(Kruskal(nodes, edges))
 


if __name__ == '__main__':
    main()

---

二、Prim算法

1.数据结构：邻接矩阵

2.思路

（1）初始化visit、dist列表

（2）下述步骤重复n次

• 找出dist列表中最小的距离以及其下标nextIndex
• 通过graph中nextIndex行的元素更新dist

3.代码实现

import math

def prim(graph, start=0):

    visit_list = []                # 记录每次选择的顶点
    vertex_num = len(graph)
    visit = [False] * vertex_num
    dist = [math.inf] * vertex_num # 选择某个顶点的边长度

    for i in range(vertex_num):

        minDist = math.inf
        nextIndex = start
        # 寻找dis列表中最小的距离以及下标
        for j in range(vertex_num):
            if dist[j] < minDist and not visit[j]:
                minDist = dist[j]
                nextIndex = j

        visit_list.append(nextIndex)
        visit[nextIndex] = True

        # 通过graph中nextIndex行的元素更新dist
        for j in range(vertex_num):
            if dist[j] > graph[nextIndex][j] and not visit[j]:
                dist[j] = graph[nextIndex][j]
    return dist, visit_list

def main():
    # 邻接矩阵存储图
    graph = [
        [0, 6, 3, math.inf, math.inf, math.inf],
        [6, 0, 2, 5, math.inf, math.inf],
        [3, 2, 0, 3, 4, math.inf],
        [math.inf, 5, 3, 0, 2, 3],
        [math.inf, math.inf, 4, 2, 0, 5],
        [math.inf, math.inf, math.inf, 3, 5, 0],
    ]

    dis, visit = prim(graph) # 最小生成树边的集合:dis[:start] union dis[start+1:]
    print(dis, visit)  # [inf, 2, 3, 3, 2, 3] [0, 2, 1, 3, 4, 5]

if __name__ == '__main__':
    main()