图 - 最小生成树

#!/usr/bin/env python3.3
# -*- coding:utf-8 -*-
# Copyright 2013

'''
无向图最小生成树：
1）kruskal算法：将所有的边按权值从小到大排序，顺序遍历把两个顶点不在同一个连通分量的边添加到最小生成树
2）Prim算法：从单一顶点开始，逐条添加连接最小生成树与一个不在树中的顶点的最小权边
'''

def mst_kruskal(graph, graph_size):
    # 并查集
    ufs = UnionFindSet()
    for v in range(graph_size):
        ufs.make_set(v)
    # 将所有的边按权值从小到大排序
    edge_queue = []
    for u in range(graph_size):
        for v in range(graph_size):
            if graph[u][v] is not None:
                edge_queue.append((graph[u][v], u, v))
    edge_queue.sort()

    mst = []
    mst_weight = 0
    for edge in edge_queue:
        # 判断两个顶点是否在同一个连通分量中
        weight, u, v = edge
        if ufs.find(u) != ufs.find(v):
            # 把边加入到最小生成树
            ufs.union(u, v)
            mst.append((u, v))
            mst_weight += weight
    return mst, mst_weight

def mst_prim(graph, graph_size):
    visit_set = {0}
    low_weight_from = [0] * graph_size
    low_weight = [w for w in graph[0]]
    low_weight[0] = 0

    mst = []
    mst_weight = 0
    for repeat in range(1, graph_size):
        # 查找最小权值顶点
        k = None
        for v in range(graph_size):
            if v not in visit_set and low_weight[v] is not None:
                if k is None or low_weight[v] < low_weight[k]:
                    k = v
        # 更新相邻顶点权值
        for v in range(graph_size):
            if v not in visit_set and graph[k][v] is not None:
                if low_weight[v] is None or graph[k][v] < low_weight[v]:
                    low_weight_from[v] = k
                    low_weight[v] = graph[k][v]
        # 把边加入到最小生成树
        mst.append((low_weight_from[k], k))
        mst_weight += low_weight[k]
        visit_set.add(k)
    return mst, mst_weight