题目链接、描述
https://www.lintcode.com/problem/3672
在本题中共存在 n 个城市,其编号范围为 1 到 n。
同时存在一个 connections 数组且 connections[i] = [a_i, b_i, c_i]connections[i]=[a_i
,b_i ,c_i ],其含义为联通城市 a_i 和 b_i 的成本为 c_i
请返回连通全部城市所需要的最低成本。如果无法联通全部城市,返回 -1。
样例
样例 1
输入:
3
[[1,2,1], [2,3,2], [1,3,3]]
输出:
3
解释:
选择 [1,2,1] 和 [2,3,2] 即可联通全部 n 个城市,此时花费最少,为 3。
样例 2
输入:
3
[[1,2,1]]
输出:
-1
解释:
无法根据 connections 联通所有的城市。
思路
并查集
答案
public class Solution {
/**
* @param n: the number of cities
* @param connections: the connection info between cities
* @return: nothing
*/
public int minimumCost(int n, int[][] connections) {
Arrays.sort(connections, new Comparator<int[]>() {
@Override
public int compare(int[] a, int[] b) {
return a[2] - b[2];
}
});
UF uf = new UF(n);
int ans = 0;
for (int[] arr : connections) {
int a = arr[0], b = arr[1], cost = arr[2];
int f1 = uf.find(a);
int f2 = uf.find(b);
if (f1 == f2) continue;
ans += cost;
uf.union(a, b);
if (uf.sets == 1) return ans;
}
return -1;
}
static class UF {
int[] parents;
int[] size;
int[] help;
int sets;
public UF(int n) {
sets = n;
parents = new int[n + 1];
size = new int[n + 1];
help = new int[n + 1];
for (int i = 1; i <= n; i++) {
parents[i] = i;
size[i] = 1;
}
}
public int find(int x) {
int hi = 0;
while (x != parents[x]) {
help[hi++] = x;
x = parents[x];
}
for (hi--; hi >= 0; hi--) {
parents[help[hi]] = x;
}
return x;
}
public void union(int a, int b) {
int f1 = find(a);
int f2 = find(b);
if (f1 != f2) {
if (size[f1] >= size[f2]) {
size[f1] += size[f2];
parents[f2] = f1;
} else {
size[f2] += size[f1];
parents[f1] = f2;
}
sets--;
}
}
}
}