Union find with path compression (Java)

class UnionFind {
    int count;
    int[] parent;
    int[] size;

    public UnionFind(int n) {
        parent = new int[n];
        size = new int[n];
        count = n;
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            size[i] = 1;
        }
    }

    public find(int a) {
        int root = a;
        while (parent[root] != root)
            root = parent[root];
        compress(a, root);
        return root;
    }

    public void compress(int a, int root) {
        while (a != root) {
            int next = parent[a];
            parent[a] = root;
            a = next;
        }
    }

    public void union(int a, int b) {

        if (find(a) == find(b))
            return;

        int ap = parent[a];
        int bp = parent[b];

        if (size[ap] <= size[bp]) {
            parent[ap] = bp;
            size[bp] += size[ap];
            compress(a, bp);
        } else {
            parent[bp] = ap;
            size[ap] += size[bp];
            compress(b, ap);
        }
        count--;
    }
}