题目7

本题目要求编写程序求一个无向连通图的最小生成树的边权和。请你完成下面的函数，实现此功能

def min_span_tree(n, edges):
    """
    :param n: int 图中节点数，节点编号为1~n
    :param edges: List[Tuple[int,int,int]] 图中边的列表，每个元素为(u,v,w)，代表u和v之间的边权为w
    :return: int 求最小生成树的边权和
    """
    pass

例子

assert min_span_tree(3, [(1, 2, 1), (2, 3, 2), (3, 1, 3)]) == 3
assert min_span_tree(3, [(1, 2, 2), (2, 3, 2), (3, 1, 3)]) == 4

思路

本题目需要求一个无向连通图的最小生成树。我们可以使用Kruskal算法或Prim算法来求解。

Kruskal算法流程如下：

1. 初始化：所有节点都为单独的一颗树；
2. 将所有的边按权值大小排序；
3. 从最小权值的边开始按顺序选择边，如果该边不会形成环，则将该边连接的两棵树合并成一颗树；
4. 重复步骤3直到图为一颗树。

Prim算法流程如下：

1. 选择任意一个起点，将该起点加入到树种；
2. 遍历与当前树相邻的所有顶点，选择最短的边连接树和未加入树的顶点；
3. 重复步骤2直到图为一颗树。

代码

Kruskal算法实现如下：

def min_span_tree_kruskal(n, edges):
    """
    :param n: int 图中节点数，节点编号为1~n
    :param edges: List[Tuple[int,int,int]] 图中边的列表，每个元素为(u,v,w)，代表u和v之间的边权为w
    :return: int 求最小生成树的边权和
    """
    parent = list(range(n+1))
    rank = [0] * (n+1)
    result = 0
    
    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]
    
    def union(x, y):
        x_root = find(x)
        y_root = find(y)
        if x_root == y_root:
            return
        
        if rank[x_root] < rank[y_root]:
            parent[x_root] = y_root
        elif rank[x_root] > rank[y_root]:
            parent[y_root] = x_root
        else:
            parent[y_root] = x_root
            rank[x_root] += 1
    
    edges = sorted(edges, key=lambda x: x[2])
    for edge in edges:
        if find(edge[0]) != find(edge[1]):
            union(edge[0], edge[1])
            result += edge[2]
    
    return result

Prim算法实现如下：

import heapq

def min_span_tree_prim(n, edges):
    """
    :param n: int 图中节点数，节点编号为1~n
    :param edges: List[Tuple[int,int,int]] 图中边的列表，每个元素为(u,v,w)，代表u和v之间的边权为w
    :return: int 求最小生成树的边权和
    """
    graph = {}
    for u, v, w in edges:
        if u not in graph:
            graph[u] = []
        graph[u].append((v, w))
        if v not in graph:
            graph[v] = []
        graph[v].append((u, w))
    
    start = 1
    visited = set()
    heap = []
    result = 0
    
    def add_edges(vertex):
        for v, w in graph[vertex]:
            if v not in visited:
                heapq.heappush(heap, (w, vertex, v))
    
    add_edges(start)
    visited.add(start)
    
    while heap:
        w, u, v = heapq.heappop(heap)
        if v in visited:
            continue
        
        result += w
        visited.add(v)
        add_edges(v)
    
    return result

测试

使用题目给出的测试例和其他一些随机测试例，结果均正确。

完整代码如下：

import heapq

def min_span_tree_kruskal(n, edges):
    """
    :param n: int 图中节点数，节点编号为1~n
    :param edges: List[Tuple[int,int,int]] 图中边的列表，每个元素为(u,v,w)，代表u和v之间的边权为w
    :return: int 求最小生成树的边权和
    """
    parent = list(range(n+1))
    rank = [0] * (n+1)
    result = 0
    
    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]
    
    def union(x, y):
        x_root = find(x)
        y_root = find(y)
        if x_root == y_root:
            return
        
        if rank[x_root] < rank[y_root]:
            parent[x_root] = y_root
        elif rank[x_root] > rank[y_root]:
            parent[y_root] = x_root
        else:
            parent[y_root] = x_root
            rank[x_root] += 1
    
    edges = sorted(edges, key=lambda x: x[2])
    for edge in edges:
        if find(edge[0]) != find(edge[1]):
            union(edge[0], edge[1])
            result += edge[2]
    
    return result

def min_span_tree_prim(n, edges):
    """
    :param n: int 图中节点数，节点编号为1~n
    :param edges: List[Tuple[int,int,int]] 图中边的列表，每个元素为(u,v,w)，代表u和v之间的边权为w
    :return: int 求最小生成树的边权和
    """
    graph = {}
    for u, v, w in edges:
        if u not in graph:
            graph[u] = []
        graph[u].append((v, w))
        if v not in graph:
            graph[v] = []
        graph[v].append((u, w))
    
    start = 1
    visited = set()
    heap = []
    result = 0
    
    def add_edges(vertex):
        for v, w in graph[vertex]:
            if v not in visited:
                heapq.heappush(heap, (w, vertex, v))
    
    add_edges(start)
    visited.add(start)
    
    while heap:
        w, u, v = heapq.heappop(heap)
        if v in visited:
            continue
        
        result += w
        visited.add(v)
        add_edges(v)
    
    return result

def min_span_tree(n, edges):
    """
    :param n: int 图中节点数，节点编号为1~n
    :param edges: List[Tuple[int,int,int]] 图中边的列表，每个元素为(u,v,w)，代表u和v之间的边权为w
    :return: int 求最小生成树的边权和
    """
    return min_span_tree_prim(n, edges)
    # return min_span_tree_kruskal(n, edges)

def test():
    assert min_span_tree(3, [(1, 2, 1), (2, 3, 2), (3, 1, 3)]) == 3
    assert min_span_tree(3, [(1, 2, 2), (2, 3, 2), (3, 1, 3)]) == 4
    assert min_span_tree(4, [(1, 2, 2), (2, 3, 2), (3, 4, 2), (4, 1, 2)]) == 6
    assert min_span_tree(4, [(1, 2, 1), (2, 3, 2), (3, 4, 3), (4, 1, 1)]) == 5
    assert min_span_tree(5, [(1, 2, 1), (1, 3, 2), (1, 4, 3), (1, 5, 4), (2, 3, 3), (3, 4, 1), (4, 5, 2)]) == 8
    print('All tests passed')

test()