并查集

Created in April 16, 2023

2023   ·   code   ·   LeetCode

并查集(union-find set)是一种用于管理元素所属集合的数据结构。它的实现为一个森林，其中每棵树表示一个集合，树中的节点表示对应集合中的元素。

并查集需要实现合并(union)和查询(find)两个操作：

  • 合并：合并两个元素所属集合(合并对应的树)
  • 查询：查询某个元素所属的集合(树的根节点)，一般可以用来判断两个元素是否来自同一集合

并查集的经典应用在于处理动态连接(dynamic connectivity)问题，即判断两个对象是否相互联通或者图上有多少不连通的集合。因此在初始化并查集时我们需要初始化两个属性：self.pa用来记录每个节点的父节点，以及self.count用来记录图上的集合数。



1
2
3
4

 
class UnionFindSet:
    def __init__(self, n: int):
        self.pa    = [i for i in range(n)]  ## 初始化每个节点的父节点为其自身
        self.count = n                      ## 初始化集合数为全部节点数

使用并查集进行查询时需要调用find(x)方法返回x节点的根节点，这里还可以结合路径压缩的技巧来加快查询速度。



1
2
3
4
5

 
def find(self, x: int) -> int:
    if self.pa[x] != x:
        self.pa[x] = self.find(self.pa[x])  ## 路径压缩
    
    return self.pa[x]

在find(x)的基础上我们可以实现union(x, y)方法用来合并x和y两个节点所在集合。



1
2
3
4
5
6
7

 
def union(self, x: int, y: int) -> None:
    root_x = self.find(x)                   ## 查询x的根节点
    root_y = self.find(y)                   ## 查询y的根节点

    if root_x != root_y:
        self.pa[root_y] = root_x            ## 当x和y来自不同集合时把y的根节点连到x根节点上
        self.count -= 1                     ## 同时集合数量减1

并查集的标准实现可以参考如下。



1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

 
class UnionFindSet:
    def __init__(self, n: int):
        self.pa    = [i for i in range(n)]
        self.count = n
    
    def union(self, x: int, y: int) -> None:
        root_x = self.find(x)
        root_y = self.find(y)

        if root_x != root_y:
            self.pa[root_y] = root_x
            self.count -= 1
    
    def find(self, x: int) -> int:
        if self.pa[x] != x:
            self.pa[x] = self.find(self.pa[x])
        
        return self.pa[x]

冗余连接

684. 冗余连接

树可以看成是一个连通且无环的无向图。

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

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

示例1：



1
2

 
输入：edges = [[1,2], [1,3], [2,3]]
输出：[2,3]

示例2：



1
2

 
输入：edges = [[1,2], [2,3], [3,4], [1,4], [1,5]]
输出：[1,4]

提示：

  • n == edges.length
  • 3 <= n <= 1000
  • edges[i].length == 2
  • 1 <= ai < bi <= edges.length
  • ai != bi
  • edges中无重复元素
  • 给定的图是连通的

Solution

本题的解法在于对edges进行遍历，如果边的两个端点u和v位于不同集合则将它们合并，否则说明它们已经连接到一起了直接返回即可。

题目链接：

python代码：



1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

 
class UnionFindSet:
    def __init__(self, n: int):
        self.pa    = [i for i in range(n)]
        self.count = n
    
    def union(self, x, y) -> None:
        root_x = self.find(x)
        root_y = self.find(y)

        if root_x != root_y:
            self.pa[root_y] = root_x
            self.count -= 1
    
    def find(self, x: int) -> int:
        if self.pa[x] != x:
            self.pa[x] = self.find(self.pa[x])
        
        return self.pa[x]

class Solution:
    def findRedundantConnection(self, edges: List[List[int]]) -> List[int]:
        n = len(edges)

        UF = UnionFindSet(n+1)

        for u, v in edges:
            if UF.find(u) == UF.find(v):
                return [u, v]
            else:
                UF.union(u, v)
        
        return []

C++代码：



1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

 
class Solution {
public:
    vector<int> findRedundantConnection(vector<vector<int>>& edges) {
        int N = edges.size();
        init(N+1);

        for (int i=0; i<N; ++i) {
            int u = edges[i][0], v = edges[i][1];

            if (find(u) == find(v)) return edges[i];
            else join(u, v);
        }

        return {};
    }

private:
    vector<int> pa;

    void init(int N) {
        pa.clear();
        pa.resize(N);

        for (int i=0; i<N; ++i) pa[i] = i;
    }

    void join(int u, int v) {
        u = find(u);
        v = find(v);

        if (u != v) pa[v] = u;
    }

    int find(int u) {
        if (pa[u] != u) pa[u] = find(pa[u]);

        return pa[u];
    }
};

685. 冗余连接 II

在本问题中，有根树指满足以下条件的有向图。该树只有一个根节点，所有其他节点都是该根节点的后继。该树除了根节点之外的每一个节点都有且只有一个父节点，而根节点没有父节点。

输入一个有向图，该图由一个有着n个节点(节点值不重复，从1到n)的树及一条附加的有向边构成。附加的边包含在1到n中的两个不同顶点间，这条附加的边不属于树中已存在的边。

结果图是一个以边组成的二维数组edges。每个元素是一对[uᵢ, vᵢ]，用以表示有向图中连接顶点uᵢ和顶点vᵢ的边，其中uᵢ是vᵢ的一个父节点。

返回一条能删除的边，使得剩下的图是有n个节点的有根树。若有多个答案，返回最后出现在给定二维数组的答案。

示例1：



1
2

 
输入：edges = [[1,2],[1,3],[2,3]]
输出：[2,3]

示例2：



1
2

 
输入：edges = [[1,2],[2,3],[3,4],[4,1],[1,5]]
输出：[4,1]

提示：

  • n == edges.length
  • 3 <= n <= 1000
  • edges[i].length == 2
  • 1 <= uᵢ, vᵢ <= n

Solution

本题的整体思路类似于冗余连接，我们需要对edges进行遍历并尝试构造出树结构。由于整个图上只有一条额外的边，这条边有三种可能性：

  1. 这条边破坏了树结构，使得某个节点v有两个父节点(如示例1中的节点3)
  2. 这条边导致图上出现了环(如示例2中的[4, 1])
  3. 这条边既破坏了树结构又产生了环

对于树结构，我们设置一个数组parent用来记录每个节点当前的父节点。我们把parent初始化为节点自身的编号，表示此时没有添加任何边，然后在对边进行遍历时检查parent[v]是否与v相等。如果发现parent[v] != v说明v节点已经存在一个父节点了，此时使用变量conflict来记录这条边的编号；而如果发现parent[v] == v则说明v节点还没有任何父节点，令parent[v] = u同时利用并查集检查u和v两个节点是否联通，如果联通则还需要利用变量cycle记录下这条边的编号，最后合并u和v这两个节点。

完成遍历后我们需要结合conflict和cycle两个变量进行讨论：

  • 如果conflict < 0说明图上只存在环，此时可以直接返回edges[cycle]
  • 如果conflict >=0说明图上有节点存在两个父节点，需要进一步讨论：
    • 如果cycle < 0说明图上没有环，直接返回edges[conflict]
    • 如果cycle >=0说明图上存在环，我们需要先提取具有两个父节点的子节点u, v = edges[conflict]然后返回v节点当前的边[parent[v], v]

整个算法流程可以参考如下。

题目链接：

python代码：



1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

 
class UnionFindSet:
    def __init__(self, n: int):
        self.pa    = [i for i in range(n)]
        self.count = n
    
    def find(self, x: int) -> int:
        if self.pa[x] != x:
            self.pa[x] = self.find(self.pa[x])
        return self.pa[x]
    
    def union(self, x: int, y: int) -> None:
        root_x = self.find(x)
        root_y = self.find(y)

        if root_x != root_y:
            self.pa[root_y] = root_x
            self.count -= 1

class Solution:
    def findRedundantDirectedConnection(self, edges: List[List[int]]) -> List[int]:
        n = len(edges)

        UF     = UnionFindSet(n+1)
        parent = [i for i in range(n+1)]

        conflict = -1
        cycle    = -1

        for i, (u, v) in enumerate(edges):
            if parent[v] == v:
                parent[v] = u

                if UF.find(u) == UF.find(v):
                    cycle = i
                
                UF.union(u, v)
            
            else:
                conflict = i
        
        if conflict < 0:
            return edges[cycle]
        else:
            if cycle < 0:
                return edges[conflict]
            else:
                u, v = edges[conflict]
                return [parent[v], v]

C++代码：



1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

 
struct UnionFind {
    vector<int> pa;

    UnionFind(int N) {
        pa.clear();
        pa.resize(N);

        for (int i=0; i<N; ++i) pa[i] = i;
    }

    void join(int u, int v) {
        u = find(u);
        v = find(v);

        if (u != v) pa[v] = u;
    }

    int find(int u) {
        if (pa[u] != u) pa[u] = find(pa[u]);

        return pa[u];
    }
};

class Solution {
public:
    vector<int> findRedundantDirectedConnection(vector<vector<int>>& edges) {
        int N = edges.size();
        UnionFind UF(N+1);

        vector<int> parent(N+1);
        for (int i=0; i<N+1; ++i) parent[i] = i;

        int conflict = -1;
        int cycle    = -1;

        for (int i=0; i<N; ++i) {
            int u = edges[i][0];
            int v = edges[i][1];

            if (parent[v] != v) {
                conflict = i;
            } else {
                parent[v] = u;

                if (UF.find(u) == UF.find(v)) cycle = i;
                else UF.join(u, v);
            }
        }

        if (conflict < 0) return {edges[cycle][0], edges[cycle][1]};
        else {
            if (cycle < 0) return {edges[conflict][0], edges[conflict][1]};
            else return {parent[edges[conflict][1]], edges[conflict][1]};
        }
    }
};