算法学习（6）：LeetCode刷题之并查集

前言

并查集是一种巧妙的算法思想，主要解决连通性问题的场景，比如：一组对象被划分成了若干个区域，求解被划分区域的数量，或者求解如何将这几个区域连接起来。

题目中往往以图形的形式来表示区域，图中节点的连通关系用一个二维数组来表示，就比如LeetCode No. 1319题，题目给出了图中节点的数量，还给出了一个二维数组，如下输入表示有4个节点，其中0和1、0和2、1和2之间有连接。

输入：n = 4, connections = [[0,1],[0,2],[1,2]]

我们怎么判断2个节点之间是否连通呢？我们可以把每个区域看作一棵树，每个区域推举一个节点作为根节点root，如果2个节点分别找到的root节点是一样的，说明这2个节点在同一个区域，即连通的。反之，如果找到的root节点不一样，表示这2个节点属于不同的区域，即不连通。

那么问题来了，我们怎么简单的表示这种树形结构呢，答案是用一个数组，下标表示当前节点，下标对应的数组值表示当前节点的父节点，如果是根节点，就指向它自己。举个例子：

上面的2棵树分别是两个不同区域中的节点，每个区域推举出一个节点作为root根节点，我们要判断2个节点是否连通，就判断它们的root节点是否是同一个。

说完并查集的使用场景，我们来详细介绍下并查集的算法思路。并查集的英文名字是Union Find，代表了并查集这种数据结构中最重要的2个操作：find()函数用来找到当前节点的根节点，union()函数用来将2个节点进行连接。下面来编码实现并查集：

class UnionFind {
    private int[] connect; // 一维数组记录每个节点的父节点
    private int count; // 联通分量的数量，即树的个数，区域的个数

    // 构造方法
    public UnionFind(int n) {
        // n是元素个数
        connect = new int[n];
        // 初始化让每个节点的父节点指向自己，下面的union()方法会改变节点间的连通关系
        for (int i = 0; i < n; i++) {
            connect[i] = i;
        }
        // 初始化区域的个数就是节点的个数
        count = n;
    }

    // find方法用来找到当前节点的最顶层节点
    public int find(int n) {
        // 迭代方式
        while (n != connect[n]) {
            n = connect[n];
        }
        return n;
    }

    // union方法用来连接2个非连通域
    public void union(int m, int n) {
        // 首先找到2个节点的root节点
        int root1 = find(m);
        int root2 = find(n);
        // 如果是相同的root，直接返回
        if (root1 == root2) {
            return;
        }
        // 连接很简单，只需将一个root节点指向另一个root节点
        connect[root1] = root2;
        // 连通了2个区域后，区域总数减1
        count--;
    }

    // 判断2个节点是否联通
    public boolean isConnect(int m, int n) {
        return find(m) == find(n);
    }

    // 返回连通分量的个数，即树的个数
    public int getCount() {
        return count;
    }
}

上面的find()函数使用了迭代方式，还可以使用递归方式

// find方法用来找到当前节点的最顶层节点
public int find(int n) {
    // 递归方式
    if (n == connect[n]) {
        return n;
    }
    return find(connect[n]);
}

上面代码中维护一个count变量很有用，它用来记录图中有多少个连通分量，能够应用于不同题目的场景。

我们再来看算法的时间复杂度，find()方法在最坏的情况要遍历整个数组，平均时间复杂度是O(N)，union()方法和isConnect()方法依赖find()方法，从而导致这2个方法的时间复杂度也是O(N)。这个时间复杂度是很糟糕的，试想如果要解决社交网络这种规模的数据，线性时间是完全不可接受的。那么怎么来优化呢？

答案就是路径压缩和按秩合并：
路径压缩：类似于平衡二叉树的思想，为了防止二叉树退化成链表，尽可能将树保持平衡。最终可压缩成一个2层的树，连通分量中所有节点都和根节点root直接相连，如下图：

这样一来，find()函数就能在O(1)时间复杂度内找到root，相应的，union()方法和isConnect()方法也下降至O(1)。看下如何实现find()方法：

private int find(int x) {
    while (parent[x] != x) {
        // 进行路径压缩
        parent[x] = parent[parent[x]];
        x = parent[x];
    }
    return x;
}

按秩合并：之前我们只是将其中一棵树简单粗暴的接到另一棵树上，这样也会造成树不平衡，我们希望小一些的树接到大一些的树的下面，这样就更平衡一些。“秩”可以是树中节点的个数，解决方法是额外维护一个数组，记录每棵树包含的节点数量。

我们来看优化后的并查集代码：

class FastUnionFind {
    // 连通分量个数
    private int count;
    // 存储一棵树
    private int[] parent;
    // 记录树的“重量”
    private int[] size;

    public FastUnionFind(int n) {
        this.count = n;
        parent = new int[n];
        size = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            size[i] = 1;
        }
    }

    public void union(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        if (rootP == rootQ)
            return;

        // 小树接到大树下面，较平衡
        if (size[rootP] > size[rootQ]) {
            parent[rootQ] = rootP;
            size[rootP] += size[rootQ];
        } else {
            parent[rootP] = rootQ;
            size[rootQ] += size[rootP];
        }
        count--;
    }

    public boolean isConnect(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        return rootP == rootQ;
    }

    private int find(int x) {
        while (parent[x] != x) {
            // 进行路径压缩
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    }

    public int getCount() {
        return count;
    }
}

正文

1、LeetCode No. 1319 连通网络的操作次数

用以太网线缆将 n 台计算机连接成一个网络，计算机的编号从 0 到 n-1。线缆用 connections 表示，其中 connections[i] = [a, b] 连接了计算机 a 和 b。
网络中的任何一台计算机都可以通过网络直接或者间接访问同一个网络中其他任意一台计算机。
给你这个计算机网络的初始布线 connections，你可以拔开任意两台直连计算机之间的线缆，并用它连接一对未直连的计算机。请你计算并返回使所有计算机都连通所需的最少操作次数。如果不可能，则返回 -1 。 

示例 1：
输入：n = 4, connections = [[0,1],[0,2],[1,2]]
输出：1
解释：拔下计算机 1 和 2 之间的线缆，并将它插到计算机 1 和 3 上。

这道题就是典型的计算连通分量的题型，如果能够找到图中有多少个连通分量，那么要计算再接几根线才能将整个区域连通就非常简单了。使用上面的模板，

class Solution {
    public int makeConnected(int n, int[][] connections) {
    	// 如果线的数量比节点数-1还要小，肯定连不通
        if (connections.length < n - 1) {
            return -1;
        }
        // 初始化并查集结构，数组大小即是节点数量
        UnionFind unionFind = new UnionFind(n);
        // 将两两节点依次连通起来
        for (int[] conn : connections) {
            unionFind.union(conn[0], conn[1]);
        }
        // 最后剩下几个连通分量，-1就是需要连接的线条数
        return unionFind.getCount() - 1;
    }
}

class UnionFind {
	// 将上面的并查集的数据结构copy过来
}

2、LeetCode No. 547 省份数量

有 n 个城市，其中一些彼此相连，另一些没有相连。如果城市 a 与城市 b 直接相连，且城市 b 与城市 c 直接相连，那么城市 a 与城市 c 间接相连。省份 是一组直接或间接相连的城市，组内不含其他没有相连的城市。
给你一个 n x n 的矩阵 isConnected ，其中 isConnected[i][j] = 1 表示第 i 个城市和第 j 个城市直接相连，而 isConnected[i][j] = 0 表示二者不直接相连。
返回矩阵中 省份 的数量。

这个题目跟上面那个没什么本质的差别，同样是找图中的连通分量。套路是一样的，先初始化并查集，再挨个做Union操作，最后返回连通区域即可。

public int findCircleNum(int[][] isConnected) {
    UnionFind unionFind = new UnionFind(isConnected.length);
    for (int i = 0; i < isConnected.length; i++) {
        for (int j = 0; j < isConnected[i].length; j++) {
            if (isConnected[i][j] == 1) {
                unionFind.union(i, j);
            }
        }
    }
    return unionFind.getCount();
}

class UnionFind {
	// 将上面的并查集的数据结构copy过来
}

3、LeetCode No. 684 冗余连接

树可以看成是一个连通且 无环 的 无向 图。给定往一棵 n 个节点 (节点值 1～n) 的树中添加一条边后的图。添加的边的两个顶点包含在 1 到 n 中间，且这条附加的边不属于树中已存在的边。图的信息记录于长度为 n 的二维数组 edges ，edges[i] = [ai, bi] 表示图中在 ai 和 bi 之间存在一条边。

请找出一条可以删去的边，删除后可使得剩余部分是一个有着 n 个节点的树。如果有多个答案，则返回数组 edges 中最后出现的边。

这道题也是典型的求解连通问题的，可以使用一样的模板。按照题目给的数组的顺序依次判断是否是连通的，如果不连通，就让它们连通，碰到的第一个连通的数组就是冗余的。由于本题节点编码是从1到n的，而不是从0到n-1的，所以并查集数据结构中的数组大小应该是n+1，其他均没有什么变化。

public int[] findRedundantConnection(int[][] edges) {
    UnionFind unionFind = new UnionFind(edges.length);
    for (int[] tmp : edges) {
        if (!unionFind.isConnect(tmp[0], tmp[1])) {
            unionFind.union(tmp[0], tmp[1]);
        } else {
            return tmp;
        }
    }
    return new int[2];
}

class UnionFind {
	// 构造方法
    public UnionFind(int n) {
        // n是元素个数
        connect = new int[n + 1];
	// 将上面的并查集的数据结构copy过来
}

总结

并查集巧妙地将图节点的连通性问题转换成数组来表示，大大简化了操作。