The disjoint-set union (DSU) data structure maintains a partition of the integers from 0 through n. Each set is represented by a root, and find(x) returns the representative of the set containing x. A union operation merges two sets; a connectivity query checks whether two elements have the same representative.
The parent array forms a forest: a root is exactly an element whose parent is itself, and every element in a set reaches its representative by following parent links. find uses iterative path halving, redirecting each visited node to its grandparent. union attaches the smaller tree below the larger one. Together, path compression and union by size give amortized O(α(N)) time per operation, where α is the inverse Ackermann function, and O(N) space.
Operation type 0 merges the sets containing a and b; every other type asks whether they are connected. Input is read as whitespace-separated integers, so line spacing does not affect parsing.
Java
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
88
89
90
91
92
93
import java.io.BufferedInputStream;
import java.io.IOException;
public class Main {
public static void main(String[] args) throws IOException {
FastScanner input = new FastScanner();
int n = input.nextInt();
int m = input.nextInt();
DisjointSet sets = new DisjointSet(n + 1);
StringBuilder output = new StringBuilder();
for (int i = 0; i < m; i++) {
int type = input.nextInt();
int a = input.nextInt();
int b = input.nextInt();
if (type == 0) {
sets.union(a, b);
} else {
output.append(sets.find(a) == sets.find(b) ? "YES\n" : "NO\n");
}
}
System.out.print(output);
}
private static final class DisjointSet {
private final int[] parent;
private final int[] size;
DisjointSet(int count) {
parent = new int[count];
size = new int[count];
for (int i = 0; i < count; i++) {
parent[i] = i;
size[i] = 1;
}
}
int find(int element) {
while (element != parent[element]) {
parent[element] = parent[parent[element]];
element = parent[element];
}
return element;
}
void union(int a, int b) {
int rootA = find(a);
int rootB = find(b);
if (rootA == rootB) {
return;
}
if (size[rootA] < size[rootB]) {
int temporary = rootA;
rootA = rootB;
rootB = temporary;
}
parent[rootB] = rootA;
size[rootA] += size[rootB];
}
}
private static final class FastScanner {
private final BufferedInputStream input = new BufferedInputStream(System.in);
private final byte[] buffer = new byte[1 << 16];
private int length;
private int position;
private int read() throws IOException {
if (position == length) {
length = input.read(buffer);
position = 0;
if (length == -1) {
return -1;
}
}
return buffer[position++];
}
int nextInt() throws IOException {
int c;
do {
c = read();
} while (c <= ' ' && c != -1);
int value = 0;
while (c > ' ') {
value = value * 10 + c - '0';
c = read();
}
return value;
}
}
}