lintcode 3672 · 最低成本联通所有城市【中等】

题目链接、描述

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--;
            }
        }
    }


}