2.7.1 Disjoint Set Union

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

/*

Maintain a set of elements partitioned into non-overlapping subsets. Each partition is identified by
a unique representative called its root. Each subset is stored as a tree whose nodes point toward
its root: finding a representative follows parent pointers, redirecting visited nodes straight to
the root along the way (path compression), while unions attach the smaller tree beneath the root of
the larger one (union-by-size). Together these keep the trees nearly flat. Union-by-size is
interchangeable with the union-by-rank strategy used by Sparse Disjoint Set Union and yields the
same complexity bounds, but union-by-size additionally supports `set_size()` queries.

This version is simplified to work only on the fixed set of integer elements $[0, n)$, chosen at
construction. For arbitrary element types, see Sparse Disjoint Set Union.

- `DSU(n)` constructs `n` singleton partitions over elements $[0, `n`)$.
- `find(u)` returns the unique representative of the partition containing `u`.
- `sets()` returns the current number of partitions.
- `set_size(u)` returns the number of elements in the partition containing `u`.
- `is_united(u, v)` returns whether elements `u` and `v` belong to the same partition.
- `unite(u, v)` replaces the partitions containing `u` and `v` with a single new partition
  consisting of the union of elements in the original partitions, returning `true` if a merge
  occurred or `false` if `u` and `v` already belonged to the same partition.

Time Complexity:
- O(n) per call to the constructor, and O(1) per call to `sets()`.
- O(alpha(n)) per call to `find()`, `set_size()`, `is_united()`, and `unite()`, where $n$ is the
  number of elements, and $\alpha(n)$ is the extremely slow growing inverse of the Ackermann
  function (effectively a very small constant for all practical values of $n$).

Space Complexity:
- O(n) for storage of the disjoint set forest elements.
- O(log n) auxiliary stack space for operations that call `find()`.
- O(1) auxiliary for `sets()`.

*/

#include <numeric>
#include <utility>
#include <vector>

class DSU {
  std::vector<int> root, size;
  int num_sets;

 public:
  explicit DSU(int n) : root(n), size(n, 1), num_sets(n) { std::iota(root.begin(), root.end(), 0); }

  int find(int u) { return root[u] == u ? u : root[u] = find(root[u]); }
  int sets() const { return num_sets; }
  int set_size(int u) { return size[find(u)]; }
  bool is_united(int u, int v) { return find(u) == find(v); }

  bool unite(int u, int v) {
    u = find(u);
    v = find(v);
    if (u == v) {
      return false;
    }
    num_sets--;
    if (size[u] < size[v]) {
      std::swap(u, v);
    }
    root[v] = u;
    size[u] += size[v];
    return true;
  }
};

/*** Example Usage ***/

#include <cassert>

int main() {
  DSU dsu(7);
  assert(dsu.unite(0, 1));
  assert(dsu.unite(1, 5));
  assert(dsu.unite(3, 4));
  assert(dsu.unite(3, 6));
  assert(!dsu.unite(0, 5));  // Already united.
  assert(dsu.sets() == 3);
  assert(dsu.set_size(1) == 3);
  assert(dsu.set_size(2) == 1);
  assert(dsu.is_united(0, 1));
  assert(!dsu.is_united(0, 2));
  assert(!dsu.is_united(1, 6));
  assert(dsu.is_united(4, 6));
  return 0;
}