Categories
interview

Union find with path compression (Js, Java)

Javascript (ES6)

class UnionFind {
  constructor() {
    this.nodes = 0;
    this.sizeOf = [];
    this.parent = [];
    for (let i = 0; i < 10; i++) {
      this.nodes++;
      this.size[i] = 1;
      this.parent[i] = i;
    }
  }

  rootOf = (a) => {
    let current = a;
    while (this.parent[current] !== current) {
      current = this.parent[current];
    }
    this.compress(a, current);
    return current;
  };

  compress = (a, root) => {
    let curr = a;
    while (this.parent[current] !== current) {
      const next = this.parent[current];
      this.parent[current] = root;
      curr = next;
    }
    return current;
  };

  union = (a, b) => {
    const ap = rootOf(a);
    const bp = rootOf(b);
    if (this.sizeOf[ap] < this.sizeOf[bp]) {
      this.parent[ap] = bp;
      this.sizeOf[bp] += this.sizeOf[ap];
      this.compress(a, bp);
    } else {
      this.parent[bp] = ap;
      this.sizeOf[ap] += this.sizeOf[bp];
      this.compress(b, ap);
    }
    this.nodes--;
  };
}

Demo

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--;
    }
}