Union-Find并查集算法

1. 概念

Union-Find 算法通常叫做并查集算法，它主要用于处理集合的合并和查询问题。顾名思义，Union-就是合并，Find-就是查找。举个典型的例子：

这是一个图，图中共有 9 个顶点，其中有的顶点是相连的，由此可以引出以下几个问题（这里需要你了解图的连通性等概念）：

1. 怎么快速判断图中任意两个顶点是否连通？例如顶点 0 与顶点 1 是连通的，顶点 2 与顶点 4 是连通的，但是顶点 1 与 顶点 2 是不连通的。
2. 怎么计算图中的连通分量个数？一个连通分量就是一个彼此相连通的顶点的集合，例如上图中共有 3 个连通分量。
3. 怎么将任意 2 个顶点连通起来？例如将顶点 1 和顶点 2 连通起来，最终就只剩下 2 个连通分量了。

理解了这 3 个问题，基本上就理解了 Union-Find 算法了，由此我们可以定义出一个 Union-Find 算法的雏形出来。

class UnionFind(n: Int) {

    /**
     * 判断顶点 p、q 是否连通
     */
    fun isConnected(p: Int, q: Int): Boolean {
        return false        
    }

    /**
     * 连通顶点 p、q
     */
    fun union(p: Int, q: Int) {
    }

    /**
     * 连通分量的个数
     */
    fun count(): Int {
        return 0
    }   
}

2. 怎么判断 p 和 q 是连通的

最容易想到的方法，首先遍历找出所有与顶点 p 连通的顶点，如果里面包含顶点 q，说明顶点 p、q 是连通的，否则就没有连通。但是这个判断逻辑很复杂，时间复杂度也高，搞不好就是O(n)的时间复杂度。那么 Union-Find 是怎么处理的呢？它很巧妙地用了树的结构，为每个顶点都设置了一个指针，指向其父节点，如果是根节点的话，该指针指向自己。这样每个连通分量都会形成一棵树，每棵树都只有一个根节点，要判断 p、q 是否连通只需要快速找到 p、q 在所在连通分量的根节点是否相同即可。

3. 具体思路

初始时，每个顶点的父节点指针均指向自己，如下图所示：

那么接下来怎么将顶点一个一个连通起来呢，我们只需将其中一个顶点的根节点接到另一个顶点的根节点上，过程如下图所示：

所有的顶点按这种方式连接起来之后，最开始提的几个问题都可以迎刃而解了。判断 p、q 是否连通，分别找到 p 的根节点 rootP，q 的根节点 rootQ，如果 rootP == rootQ 就说明它们是连通的。计算连通分量，初始时连通分量的个数就是顶点的总个数，每连通 2 个不同连通分量的顶点时，则将连通分量数减1。

至此，我们可以初步写出 Union-Find 算法的代码了：

/**
 * 并查集
 *
 * @param n 顶点个数，为了方便我们这里从 0 开始编号
 */
class UnionFind(n: Int) {

    //连通分量的个数，初始为顶点个数
    private var count = n
    //存储每个节点的父节点编号
    private var parent = IntArray(n)

    init {
        //初始每个节点的父节点指针都指向自己
        for (i in parent.indices) {
            parent[i] = i
        }
    }

    /**
     * 找到根节点
     */
    private fun findRoot(p: Int): Int {
        var curr = p
        //根节点的父节点都指向自己
        while (curr != parent[curr]) {
            curr = parent[curr]
        }
        return curr
    }

    /**
     * 连通顶点 p、q
     */
    fun union(p: Int, q: Int) {
        //分别找到 p, q 的根节点
        var rootP = findRoot(p)
        var rootQ = findRoot(q)
        //只有不连通的两个顶点才需要连通
        if (rootP != rootQ) {
            //将其中一个节点接到另一个上，这里将 rootP 接到 rootQ 上
            parent[rootP] = rootQ
            count--
        }
    }

    /**
     * 判断顶点 p、q 是否连通
     */
    fun isConnected(p: Int, q: Int): Boolean {
        var rootP = findRoot(p)
        var rootQ = findRoot(q)
        //根节点相同就表示是连通的
        return rootP == rootQ
    }

    /**
     * 连通分量的个数
     */
    fun count(): Int {
        return count
    }

}

4. 平衡性优化

前面代码中可以看到，不管是union() 方法还是isConnected ()方法，其效率的高低都与findRoot()有直接关系，而 findRoot()又与树的高度有直接关系。如果是一棵平衡的树(注意这里不一定是二叉树了)，其时间复杂度我们可以近似认为是 O(logn)，但如果这棵树极度不平衡，近似一个线性链表，那么其时间复杂度则退化为 O(n) 了。前面合并两颗树的时候，我们随便选了一个根节点连接到另一个根节点上，如果我们每次合并的时候都恰好像下图这样操作，那么就会生成一棵线性的树：

上图中合并时，明显将节点 2 连接到节点 1 上是最优的，这样节点 1 就成为根节点，它只有 2 个子节点，树的高度为 2。假设我们知道根节点所在树的总节点数，那么我们在连接 2 棵不同树的根节点的时候，总是将树节点数少的根节点连接在节点数多的根节点上，那么数就总能保持平衡了。

解决方法就是额外使用一个数组来记录树的节点数，优化代码如下：

    //记录以该节点为根节点的树的总节点数
    private var size = IntArray(n)

    fun union(p: Int, q: Int) {
        //分别找到 p, q 的根节点
        var rootP = findRoot(p)
        var rootQ = findRoot(q)
        //只有不连通的两个顶点才需要连通
        if (rootP != rootQ) {
            //将树节点少的连接到树节点多的
            if (size[rootP] < size[rootQ]) {
                parent[rootP] = rootQ
                size[rootQ] = size[rootP] + size[rootQ]
            } else {
                parent[rootQ] = rootP
                size[rootP] = size[rootP] + size[rootQ]
            }
            count--
        }
    }

通过这样调整，树就会比较平衡，时间复杂度就能降到 O(logn) 了，即便大规模数据处理，也非常快了。

5. 查找路径压缩优化

前面查找某个节点的根节点时，我们都是依次向上查找父节点，直到找到根节点，尽管我们已经将树变得尽可能平衡了，但不管怎样总会要查找 logn 次，那么我们能否将每个节点的父节点都指向根节点呢，如果可以的话，每次查找根节点时都只需要查找 1 次就够了，那查找时间复杂度就可以降为 O(1) 了。

如上图所示，如果能将左边的树变成右边的树，那么任意节点的查找都只需 1 次，这将非常高效，哪怕是亿级规模的数据量。我们只需要在原来代码基础上加上一句就可：

    private fun findRoot(p: Int): Int {
        var curr = p
        //根节点的父节点都指向自己
        while (curr != parent[curr]) {
            //路径压缩的关键代码
            parent[curr] = parent[parent[curr]]
            curr = parent[curr]
        }
        return curr
    }

是不是感觉亮瞎了钛合金狗眼，一句代码就能起到这么大的优化，不能理解的画个示意图就明白了。

6. 完整算法代码

/**
 * 并查集
 *
 * @param n 顶点个数，为了方便我们这里从 0 开始编号
 */
class UnionFind(n: Int) {

    //连通分量的个数，初始为顶点个数
    private var count = n
    //存储每个节点的父节点编号
    private var parent = IntArray(n)
    //记录以该节点为根节点的树的总节点数
    private var size = IntArray(n)

    init {
        //初始每个节点的父节点指针都指向自己
        for (i in parent.indices) {
            parent[i] = i
            //初始都为1
            size[i] = 1
        }
    }

    /**
     * 找到根节点
     */
    private fun findRoot(p: Int): Int {
        var curr = p
        //根节点的父节点都指向自己
        while (curr != parent[curr]) {
            //路径压缩的关键代码
            parent[curr] = parent[parent[curr]]
            curr = parent[curr]
        }
        return curr
    }

    /**
     * 连通顶点 p、q
     */
    fun union(p: Int, q: Int) {
        //分别找到 p, q 的根节点
        var rootP = findRoot(p)
        var rootQ = findRoot(q)
        //只有不连通的两个顶点才需要连通
        if (rootP != rootQ) {
            //将树节点少的连接到树节点多的
            if (size[rootP] < size[rootQ]) {
                parent[rootP] = rootQ
                size[rootQ] = size[rootP] + size[rootQ]
            } else {
                parent[rootQ] = rootP
                size[rootP] = size[rootP] + size[rootQ]
            }
            count--
        }
    }

    /**
     * 判断顶点 p、q 是否连通
     */
    fun isConnected(p: Int, q: Int): Boolean {
        var rootP = findRoot(p)
        var rootQ = findRoot(q)
        //根节点相同就表示是连通的
        return rootP == rootQ
    }

    /**
     * 连通分量的个数
     */
    fun count(): Int {
        return count
    }

}