压缩路径的加权并查集 Java实现

import java.util.Arrays;

public class UF {
    private int[] id = null;
    private int[] sz = null;
    private int count;

    UF(int n) {
        id = new int[n];
        sz = new int[n];
        for (int i = 0; i < id.length; i++) {
            id[i] = i;
        }
        Arrays.fill(sz, 1);

        count = n+1;
    }

    public int find(int i) {
        while (i != id[i]) {
            id[i] = id[id[i]]; // 压缩路径
            i = id[i];
        }
        return i;
    }

    public void union(int a, int b) {
        int p = find(a);
        int q = find(b);
        if(p == q) return;
        if(sz[p] < sz[q]) {
            id[p] = q;
            sz[q] += sz[p];
        }
        else{
            id[q] = p;
            sz[p] += sz[q];
        }
        count--;
    }

    public boolean connected(int a, int b) {
        return find(a) == find(b);
    }

    public int count(){
        return count;
    }
}