Kruskal 算法求最小生成树

#include <iostream>
#include <climits> /* for INT_MAX */
#include <queue>
#include <algorithm>
using namespace std;
#include <stdlib.h>

/** 并查集. */
typedef struct ufs_t {
    int *p; /** 树的双亲表示法*/
    int size; /** 大小. */
} ufs_t;
/**
* @brief 创建并查集.
* @param[in] n 数组的容量
* @return 并查集
*/
ufs_t* ufs_create(int n) {
    ufs_t *ufs = (ufs_t*)malloc(sizeof(ufs_t));
    int i;
    ufs->p = (int*)malloc(n * sizeof(int));
    for(i = 0; i < n; i++)
        ufs->p[i] = -1;
    return ufs;
}
/**
* @brief 销毁并查集.
* @param[in] ufs 并查集
* @return 无
*/
void ufs_destroy(ufs_t *ufs) {
    free(ufs->p);
    free(ufs);
}
/**
* @brief Find 操作，带路径压缩，递归版.
* @param[in] s 并查集
* @param[in] x 要查找的元素
* @return 包含元素x 的树的根
*/
int ufs_find(ufs_t *ufs, int x) {
    if (ufs->p[x] < 0) return x; // 终止条件
        return ufs->p[x] = ufs_find(ufs, ufs->p[x]); /* 回溯时的压缩路径*/
}

/**
* @brief Union 操作，将y 并入到x 所在的集合.
* @param[in] s 并查集
* @param[in] x 一个元素
* @param[in] y 另一个元素
* @return 如果二者已经在同一集合，并失败，返回-1，否则返回0
*/
int ufs_union(ufs_t *ufs, int x, int y) {
    const int rx = ufs_find(ufs, x);
    const int ry = ufs_find(ufs, y);
    if(rx == ry) return -1;
    ufs->p[rx] += ufs->p[ry];
    ufs->p[ry] = rx;
    return 0;
}

/**
* @brief 获取元素所在的集合的大小
* @param[in] ufs 并查集
* @param[in] x 元素
* @return 元素所在的集合的大小
*/
int ufs_set_size(ufs_t *ufs, int x) {
    const int rx = ufs_find(ufs, x);
    return -ufs->p[rx];
}

const int MAX_NV = 11; /* 顶点数最大值*/
const int MAX_NE = 100; /* 最大边数*/
/** 边的权值类型. */
typedef int graph_weight_t;
/** 图的边. */
struct edge_t{
    int u; /** 顶点编号*/
    int v; /** 顶点编号*/
    graph_weight_t w; /** 权值*/
    bool operator>(const edge_t &other) const {
        return w > other.w;
    }
};
edge_t edges[MAX_NE];
/*
* @brief Kruskal 算法，堆+ 并查集.
* @param[in] edges 边的数组
* @param[in] n 边数，一定要大于或等于(顶点数-1)
* @param[in] m 顶点数
* @return MST 的边的权值之和
*/
graph_weight_t kruskal(const edge_t edges[], int n, int m) {
    graph_weight_t sum = 0;
    priority_queue<edge_t, vector<edge_t>,
        greater<edge_t> > q;
    ufs_t *s = ufs_create(MAX_NV);
    if (n < m - 1) return -1;
    /* 把所有边插入堆中*/
    for (int i = 0; i < n; i++) {
        q.push(edges[i]);
    }
    for (int i = 0; i < n; i++) {
        /* 从堆中退出最小权值边*/
        const edge_t e = q.top(); q.pop();
        /* 取两顶点所在集合的根*/
        const int u = ufs_find(s, e.u);
        const int v = ufs_find(s, e.v);
        if (u != v) { /* 不是同一集合，说明不连通*/
            ufs_union(s, u, v); /* 合并，连通成一个分量*/
            /* 输出生成树TE 的边，即此边加入TE */
            cout << (char)('A' + e.u) << " - " << (char)('A' + e.v) << endl;
            sum += e.w;
        }
    }
    ufs_destroy(s);
    return sum;
}

int main() {
    int m, n;
    /* 读取顶点数，边数目*/
    cin >> m >> n;
    /* 读取边信息*/
    for (int i = 0; i < n; i++) {
        char chx, chy;
        int w;
        cin >> chx >> chy >> w;
        edges[i].u = chx - 'A';
        edges[i].v = chy - 'A';
        edges[i].w = w;
    }
    /* 求解最小生成树*/
    cout << "Total : " << kruskal(edges, n, m) << endl;
    return 0;
}

/* test
输入数据:
7 11
A B 7
A D 5
B C 8
B D 9
B E 7
C E 5
D E 15
D F 6
E F 8
E G 9
F G 11
输出:
A - D
C - E
D - F
B - E
A - B
E - G
Total : 39
*/