数据结构：并查集

并查集

在生活中，经常会出现分组问题。比如一个班级分为多个小组，打篮球分为两方等等。在同一个组中的所有成员，就构成一个集合。对这种一个群体分为多个集合的数据结构，称为并查集。

其提供两个最核心的功能：

• 合并：将两个集合合并成一个集合
• 查询：查找两个元素是否属于一个集合

因此称为并查集。

实现一个并查集并不难，但是如果要实现一个高效的并查集，就需要一定的设计了。本博客讲解以C++实现的并查集，并且尽可能在时间与空间的利用上更加高效。

原理

谈到集合，在数据结构中如何维护一个集合？比如一个数组，一个set，一棵树等等。既然要探求一个最高效的存储方式，那么就要讨论如何最大化利用资源了。

如果使用一个数组来存储一个集合，那么每个集合都要开辟一个数组，在合并集合时，还需要发生数组的合并，此时又会有空间的开辟和销毁。

如果使用链式树存储集合，此时合并就会很方便：

红色与蓝色是两个不同的集合，合并集合时，只需要修改一个指针的指向即可。

但是链式结构也有问题，链式结构的数据是分散的，计算机每次加载节点都需要寻址，效率很低。有没有方法既可以保持树结构，又可以集中的存储所有数据？

如果你学习过堆，那么答案就呼之欲出了，其实就是使用一个数组形式的树。

如图，每个节点存储自己的父节点的下标，根节点存储自己的下标。

其可以转化为如下三个集合：

这是一种常见的并查集形式，但是还可以再优化。这种形式下根节点存储自己的下标，是不是可以把这块空间腾出来，存储该集合的元素个数？

如图，根节点存储的值变为负数，绝对值表示该集合的总元素个数。为什么根节点要变为负数？之前已经规定了：数组的元素存储自己父节点的下标，如果根节点的值为一个正整数，此时如何判断这是一个根节点还是普通节点，存储的值是集合总元素还是父节点下标？

因为数组下标没有负数，所以此时就可以通过正负数判断该节点是根节点还是普通节点：

• 负数：根节点，存储该集合元素总个数
• 正数：普通节点，存储父节点的下标

这是一个非常高效的存储结构，使用一个数组就表示了一个并查集，内含多个树结构。而多棵树在一起就构成了一个森林，其实并查集的本质就是一个森林。

但是至此还有一个问题，这个并查集只能表示整数集合，不能表示其它的string等类型，所以还需要一个map维持映射关系，将其他元素映射为数组下标。

实现

框架

为了提高可扩展性，把并查集定义为一个类模板，模板参数为并查集存储元素的类型。

template <typename T>
class UnionFindSet
{
private:
    vector<int> _ufs;
    map<T, int> _mp;
};

成员变量：

• _ufs：并查集的本体，用于维护集合的关系，也就是刚刚设计的那个数组
• _mp：一个映射关系，将存储的元素T映射到具体的数组下标int

初始化

初始化时并查集接收一个数组，里面是独立的元素，它们不构成任何集合关系。

随后要构建这些元素与下标的映射关系，即初始化_mp。另

最后，对于_ufs本体，全部初始化为-1。

因为一开始所有元素自成一个集合，都是集合的根节点，而根节点存储的是集合元素的个数的负数。每个集合只有一个元素，所以节点值初始化为-1。

构造函数：

UnionFindSet(vector<T>& source)
    : _ufs(source.size(), -1)
{
    for (int i = 0; i < source.size(); i++)
        _mp[source[i]] = i;
}

参数接受一个数组source，内部包含多个T类型元素，在初始化列表种将_ufs的大小扩大到与source一致，所有元素初始化为-1。

在函数体内部，完成对_mp的初始化，遍历source，存储(source[i], i)的映射关系。

合并

合并两个集合，就是将其中一个元素的根节点的父节点指针，指向另一个节点的根节点，如图：

上图展示了蓝色集合与绿色集合的合并操作，分为以下两步：

1. 将蓝色集合根节点的值加上绿色集合根节点的值：-4变-7
2. 将绿色集合的根节点的值变为蓝色集合根节点的下标：-3变0

既然要操作集合的根节点，自然就要先找到集合的根节点，写一个函数用于获取集合根节点：

int findRoot(T x)
{
    if (_mp.count(x) == 0)
        throw runtime_error("value does not exist"); // 值不存在

    int root = _mp[x];
    while (_ufs[root] >= 0)
    {
        root = _ufs[root];
    }

    return root;
}

首先通过_mp.count(x)判断该元素是否在并查集种，如果不在就抛出一个异常，表示值不存在。

随后通过一个循环，每次root = _ufs[root]，其中_ufs[root]是父节点的下标，这样就可以让root往父节点走，直到走到根节点，此时_ufs[root]是一个负数，最后跳出循环返回根节点。

找到根节点后，就可以完成集合的合并操作了：

void unionSet(T x1, T x2)
{
    int root1 = findRoot(x1);
    int root2 = findRoot(x2);

    if (root1 == root2)
        return;

    _ufs[root1] += _ufs[root2];
    _ufs[root2] = root1;
}

首先通过findRoot找到两个集合的根节点，如果根节点相同，说明两个元素本来就处于一个集合种，直接返回。

随后_ufs[root1] += _ufs[root2];完成了元素的加和，此时root1是新根，_ufs[root1]存储的是两个集合的元素总和的负数。

最后_ufs[root2] = root1;，修改toor2父节点，完成集合的合并。

这里还有一个优化，两个集合有两种合并方式：

如图，可以将绿色集合合并到蓝色集合下，也可以将蓝色集合合并到绿色集合下。这两种方式都是合理的，但是哪一种更好？

在集合种查找元素时，最多搜索树的高度次，树高度越低，那么搜索效率就越高。所以常把集合元素多的作为根。上图中因为蓝色集合元素个数多，所以把绿色集合合并到蓝色集合更优，也就是左边的方式。这个优化称为按秩合并。

代码优化：

void unionSet(T x1, T x2)
{
    int root1 = findRoot(x1);
    int root2 = findRoot(x2);

    if (root1 == root2)
        return;
    
    // 按秩合并
    if (_ufs[root1] < _ufs[root2])
    {
        _ufs[root1] += _ufs[root2];
        _ufs[root2] = root1;
    }
    else
    {
        _ufs[root2] += _ufs[root1];
        _ufs[root1] = root2;
    }
}

由于根节点存储的就是集合的元素个数，所以可以直接拿_ufs[root]来比较两个集合的大小。如果_ufs[root1] < _ufs[root2]，因为根节点存储的是负数，所以_ufs[root1]的绝对值更大，要把root2合并到root1。

查询

并查集的第二个核心操作是判断两个元素是否在同一个集合。这其实非常简单，只需要判断两个元素的根节点是否相同即可！

bool inSet(T x1, T x2)
{
    return findRoot(x1) == findRoot(x2);
}

获取成员

该接口的作用是，输入一个元素，取同一集合中的其它所有元素。

刚刚讲解过，判度两个元素是否在同一个集合，只需要看根节点是否相同。所以此处只需要：

1. 先获取输入的根节点root
2. 遍历整个并查集，判度根节点是否与root相同

vector<T> getMembers(T x) 
{
    vector<T> members;
    int root = findRoot(x);

    for (const auto& pair : _mp)
    {
        if (findRoot(pair.first) == root) 
            members.push_back(pair.first);
    }

    return members;
}

以上代码返回一个vector<T>，里面是与x为同一集合的所有元素。

首先root = findRoot(x)，获取x的根节点。随后通过for循环遍历_mp，findRoot(pair.first)获取元素根节点，再与root判等，如果相等说明在同一集合，此时尾插到members数组中。

路径压缩

当并查集使用久了，就会出现树高度太高的问题，但是并查集内部的树是多叉树，如下图两个集合：

这两个集合其实是同一个集合，但是很明显左边的树高度低，查询效率会高很多。所以并查集中常会做一个优化，将树高度尽可能降低，这个优化称为路径压缩。

压缩路径被实现在查找操作findRoot中，因为每次查找的时候，都会从树底往上遍历到根节点，这是完成路径压缩的最好时机。

路径压缩的算法核心是：

每次向上查找父节点时，把自己提高到与父节点的同一层

如图：

当前从节点4开始向上查找，首先找到父节点1，随后将4提升到与1的同一层。也就是中间的情况。

此时问题变成了：从1开始查找根节点。找到父节点7，随后将1提升到与7的同一层，此时就变成了最后一种情况。

最后找到根节点为0，由于0已经是根节点了，不能把7提升到根节点。

实现：

int findRoot(T x)
{
    if (_mp.count(x) == 0)
        throw runtime_error("value does not exist"); // 值不存在

    int root = _mp[x];
    while (_ufs[root] >= 0 && _ufs[_ufs[root]] >= 0)
    {
        _ufs[root] = _ufs[_ufs[root]]; // 路径压缩
    }

    if (_ufs[root] >= 0)
        root = _ufs[root];

    return root;
}

由于路径压缩要考虑爷爷节点是否存在，所以while内部有两个条件：_ufs[root] >= 0表示父节点存在，_ufs[_ufs[root]] >= 0表示爷爷节点存在。

只要父节点和爷爷节点都存在，那么就可以进行路径压缩，_ufs[root] = _ufs[_ufs[root]]，其中_ufs[root] 是当前节点的值存储的是父节点的下标，_ufs[_ufs[root]]是爷爷节点的下标。这个赋值将爷爷节点的下标赋值给自己，此时就把爷爷节点变成了父节点，完成了向上提升。

最后while循环离开的时候，有可能是因为爷爷节点不存在，此时root是根节点的某一个孩子，所以还要root = _ufs[root]往上走一层。

其它

还有一些其它的小接口，都很简单

• 当前并查集内部有多少个集合

size_t count()
{
    size_t size = 0;
    for (auto& num : _ufs)
    {
        if (num < 0)
            size++;
    }

    return size;
}

• 输入一个集合，获取该集合的元素个数

size_t size(T x)
{
    return abs(_ufs[findRoot(x)]);
}

想要知道集合元素个数，只需要找到根节点，然后返回绝对值即可。

总代码

• UnionFindSet.hpp

#pragma once
#include <iostream>
#include <vector>
#include <map>
#include <stdexcept>

using namespace std;

template <typename T>
class UnionFindSet
{
public:
    UnionFindSet(vector<T>& source)
        : _ufs(source.size(), -1)
    {
        for (int i = 0; i < source.size(); i++)
            _mp[source[i]] = i;
    }

    int findRoot(T x)
    {
        if (_mp.count(x) == 0)
            throw runtime_error("value does not exist"); // 值不存在

        int root = _mp[x];
        while (_ufs[root] >= 0 && _ufs[_ufs[root]] >= 0)
        {
            _ufs[root] = _ufs[_ufs[root]]; // 压缩路径
            root = _ufs[root];
        }

        if (_ufs[root] >= 0)
            root = _ufs[root];

        return root;
    }

    void unionSet(T x1, T x2)
    {
        int root1 = findRoot(x1);
        int root2 = findRoot(x2);

        if (root1 == root2)
            return;
        
        // 按秩合并
        if (_ufs[root1] < _ufs[root2])
        {
            _ufs[root1] += _ufs[root2];
            _ufs[root2] = root1;
        }
        else
        {
            _ufs[root2] += _ufs[root1];
            _ufs[root1] = root2;
        }
    }

    bool inSet(T x1, T x2)
    {
        return findRoot(x1) == findRoot(x2);
    }

    size_t count()
    {
        size_t size = 0;
        for (auto& num : _ufs)
        {
            if (num < 0)
                size++;
        }

        return size;
    }

    size_t size(T x)
    {
        return abs(_ufs[findRoot(x)]);
    }

    vector<T> getMembers(T x) 
    {
        vector<T> members;
        int root = findRoot(x);

        for (const auto& pair : _mp)
        {
            if (findRoot(pair.first) == root) 
                members.push_back(pair.first);
        }

        return members;
    }

private:
    vector<int> _ufs;
    map<T, int> _mp;
};

• test.cpp，测试代码

#include <iostream>
#include <string>
#include <vector>
#include "unionFindSet.hpp"

using namespace std;

int main()
{
    vector<string> stu = { "张三", "李四", "王五", "赵六", "翠花", "小龙", "小淘", "小明" };

    UnionFindSet<string> ufs(stu);

    cout << ufs.count() << endl;

    cout << ufs.inSet("张三", "翠花") << endl;

    ufs.unionSet("张三", "赵六");
    ufs.unionSet("王五", "小淘");
    ufs.unionSet("翠花", "小明");
    ufs.unionSet("翠花", "张三");

    cout << ufs.inSet("张三", "翠花") << endl;

    cout << ufs.count() << endl;
    cout << ufs.size("张三") << endl;

    auto members = ufs.getMembers("张三");

    for (auto& mem : members)
        cout << mem << "  ";

    cout << endl;

    return 0;
}