Weighted Quick-Union with Path Compression

 1/**
 2 * @author wangy
 3 * @version 1.0
 4 * @date 2023/5/25 / 08:50
 5 */
 6public class WeightedQuickUnionWithPathCompression {
 7    public static void main(String[] args) {
 8        QUWithPathCompression uf = new QUWithPathCompression(10);
 9        uf.union(4, 3);
10        uf.union(3, 8);
11        uf.union(6, 5);
12        uf.print();
13        StdOut.println(uf.find(4, 8));
14        StdOut.println(uf.find(3, 9));
15        uf.union(9, 4);
16        uf.union(2, 1);
17        uf.union(5, 0);
18        uf.print();
19        uf.union(7, 2);
20        uf.union(6, 1);
21        uf.union(7, 3);
22        uf.print();
23        StdOut.println(uf.find(2,9));
24        // additional test: tree reorganized
25        StdOut.println(uf.find(5, 7));
26        uf.print();
27    }
28}
29
30class QUWithPathCompression {
31    private int[] id;
32
33    public QUWithPathCompression(int N) {
34        id = new int[N];
35        for (int i = 0; i < N; i++) {
36            id[i] = i;
37        }
38    }
39
40
41    // two-pass solution:
42    // point node p and its ancestors to root node
43    // by adding another loop
44    public int root(int p) {
45        int ori_p = p;
46        while (id[p] != p) {
47            p = id[p];
48        }
49        while (id[ori_p] != p) {
50            int tmp = id[ori_p];
51            id[ori_p] = p;
52            ori_p = tmp;
53        }
54        return p;
55    }
56
57    // one-pass variant:
58    // make every other node in path point to its
59    // grandparent, thereby halving path length.
60    public int root2(int p) {
61        while (id[p] != p) {
62            id[p] = id[id[p]];  // one line extra code!
63            p = id[p];
64        }
65        return p;
66    }
67
68    // check p and q have the same root
69    // depth of p and q array accesses
70    public boolean find(int p, int q) {
71        return root(p) == root(q);
72    }
73
74    // change root of p to point to root of q
75    // depth of p and q array access
76    public void union(int p, int q) {
77        int rp = root(p);
78        int rq = root(q);
79        id[rp] = rq;
80    }
81
82    public void print() {
83        StdOut.println(Arrays.toString(id));
84    }
85}