Union & Find

UnionFind.java

class UnionFind {
    int[] parent;
    int[] rank;
    public UnionFind(int n) {
        parent = new int[n];
        rank = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
        }
    }

    // find the root of a
    public int find(int a) {
        // find with path compression, 从2ms提升到1ms
        while (parent[a] != a) {
            parent[a] = parent[parent[a]];
            a = parent[a];
        }
        return a;
    }

    public void union(int a, int b) {
        int rootA = find(a);
        int rootB = find(b);

        // union by rank
        if (rootA == rootB) return;
        if (rank[rootA] > rank[rootB]) {
            parent[rootB] = rootA;

        } else if (rank[rootA] < rank[rootB]) {
            parent[rootA] = rootB;
        } else {
            parent[rootB] = rootA;
            rank[rootA]++;
        }
    }
}