Python 算法基础篇之最小生成树算法： Prim 算法和 Kruskal 算法

import heapq

def prim(graph, start):
    min_spanning_tree = []
    visited = set()
    priority_queue = [(0, start)]

    while priority_queue:
        weight, node = heapq.heappop(priority_queue)
        if node not in visited:
            visited.add(node)
            min_spanning_tree.append((weight, node))

            for neighbor, neighbor_weight in graph[node].items():
                if neighbor not in visited:
                    heapq.heappush(priority_queue, (neighbor_weight, neighbor))

    return min_spanning_tree

def find(parent, node):
    if parent[node] != node:
        parent[node] = find(parent, parent[node])
    return parent[node]

def union(parent, rank, node1, node2):
    root1 = find(parent, node1)
    root2 = find(parent, node2)

    if root1 != root2:
        if rank[root1] > rank[root2]:
            parent[root2] = root1
        elif rank[root1] < rank[root2]:
            parent[root1] = root2
        else:
            parent[root2] = root1
            rank[root1] += 1

def kruskal(graph):
    min_spanning_tree = []
    edges = []

    for node, neighbors in graph.items():
        for neighbor, weight in neighbors.items():
            edges.append((weight, node, neighbor))

    edges.sort()

    parent = {node: node for node in graph}
    rank = {node: 0 for node in graph}

    for edge in edges:
        weight, node1, node2 = edge
        if find(parent, node1) != find(parent, node2):
            union(parent, rank, node1, node2)
            min_spanning_tree.append((weight, node1, node2))

    return min_spanning_tree

# 创建一个示例图
graph = {
    'A': {'B': 1, 'C': 4},
    'B': {'A': 1, 'C': 2, 'D': 5},
    'C': {'A': 4, 'B': 2, 'D': 1},
    'D': {'B': 5, 'C': 1}
}

# 使用Prim算法找到最小生成树
print("Prim算法最小生成树：", prim(graph, 'A'))

# 使用Kruskal算法找到最小生成树
print("Kruskal算法最小生成树：", kruskal(graph))

Prim算法最小生成树： [(0, 'A'), (1, 'B'), (1, 'C'), (1, 'D')]
Kruskal算法最小生成树： [(1, 'A', 'B'), (1, 'C', 'D'), (2, 'B', 'C')]