4
4
import java .util .Arrays ;
5
5
import java .util .List ;
6
6
7
+ /**
8
+ * The Union-Find data structure, also known as Disjoint Set Union (DSU),
9
+ * is a data structure that tracks a set of elements partitioned into
10
+ * disjoint (non-overlapping) subsets. It supports two main operations:
11
+ *
12
+ * 1. **Find**: Determine which subset a particular element is in.
13
+ * 2. **Union**: Join two subsets into a single subset.
14
+ *
15
+ * This implementation uses path compression in the `find` operation
16
+ * and union by rank in the `union` operation for efficiency.
17
+ */
7
18
public class UnionFind {
8
19
9
- private final int [] p ;
10
- private final int [] r ;
20
+ private final int [] p ; // Parent array
21
+ private final int [] r ; // Rank array
11
22
23
+ /**
24
+ * Initializes a Union-Find data structure with n elements.
25
+ * Each element is its own parent initially.
26
+ *
27
+ * @param n the number of elements
28
+ */
12
29
public UnionFind (int n ) {
13
30
p = new int [n ];
14
31
r = new int [n ];
@@ -18,19 +35,33 @@ public UnionFind(int n) {
18
35
}
19
36
}
20
37
38
+ /**
39
+ * Finds the root of the set containing the element i.
40
+ * Uses path compression to flatten the structure.
41
+ *
42
+ * @param i the element to find
43
+ * @return the root of the set
44
+ */
21
45
public int find (int i ) {
22
46
int parent = p [i ];
23
47
24
48
if (i == parent ) {
25
49
return i ;
26
50
}
27
51
52
+ // Path compression
28
53
final int result = find (parent );
29
54
p [i ] = result ;
30
-
31
55
return result ;
32
56
}
33
57
58
+ /**
59
+ * Unites the sets containing elements x and y.
60
+ * Uses union by rank to attach the smaller tree under the larger tree.
61
+ *
62
+ * @param x the first element
63
+ * @param y the second element
64
+ */
34
65
public void union (int x , int y ) {
35
66
int r0 = find (x );
36
67
int r1 = find (y );
@@ -39,6 +70,7 @@ public void union(int x, int y) {
39
70
return ;
40
71
}
41
72
73
+ // Union by rank
42
74
if (r [r0 ] > r [r1 ]) {
43
75
p [r1 ] = r0 ;
44
76
} else if (r [r1 ] > r [r0 ]) {
@@ -49,39 +81,24 @@ public void union(int x, int y) {
49
81
}
50
82
}
51
83
84
+ /**
85
+ * Counts the number of disjoint sets.
86
+ *
87
+ * @return the number of disjoint sets
88
+ */
52
89
public int count () {
53
90
List <Integer > parents = new ArrayList <>();
54
91
for (int i = 0 ; i < p .length ; i ++) {
55
- if (!parents .contains (find (i ))) {
56
- parents .add (find (i ));
92
+ int root = find (i );
93
+ if (!parents .contains (root )) {
94
+ parents .add (root );
57
95
}
58
96
}
59
97
return parents .size ();
60
98
}
61
99
100
+ @ Override
62
101
public String toString () {
63
102
return "p " + Arrays .toString (p ) + " r " + Arrays .toString (r ) + "\n " ;
64
103
}
65
-
66
- // Tests
67
- public static void main (String [] args ) {
68
- UnionFind uf = new UnionFind (5 );
69
- System .out .println ("init /w 5 (should print 'p [0, 1, 2, 3, 4] r [0, 0, 0, 0, 0]'):" );
70
- System .out .println (uf );
71
-
72
- uf .union (1 , 2 );
73
- System .out .println ("union 1 2 (should print 'p [0, 1, 1, 3, 4] r [0, 1, 0, 0, 0]'):" );
74
- System .out .println (uf );
75
-
76
- uf .union (3 , 4 );
77
- System .out .println ("union 3 4 (should print 'p [0, 1, 1, 3, 3] r [0, 1, 0, 1, 0]'):" );
78
- System .out .println (uf );
79
-
80
- uf .find (4 );
81
- System .out .println ("find 4 (should print 'p [0, 1, 1, 3, 3] r [0, 1, 0, 1, 0]'):" );
82
- System .out .println (uf );
83
-
84
- System .out .println ("count (should print '3'):" );
85
- System .out .println (uf .count ());
86
- }
87
104
}
0 commit comments