disjoint_set.h

#include <vector>

//Use Disjoint Set Forest implementation
//father[i]: the idx of father of i-th node
//sz[i]:how many node view i-th node as ancestor(and father), including itself
//If father[i]=i, means the set i belongs to only has 1 element (which is i)
namespace djs{
  struct DisjointSet{
    int n;
    std::vector<int> father;
    std::vector<int> sz;
    
    // In the beginning, size of every set is 1, the father is itself
    DisjointSet(int n){
      sz.resize(n,1);
      father.resize(n);
      std::iota(father.begin(),father.end(),0);
    }
    
    // Give x-th node, find the root
    // Running time: O(lg n)
    int find(int x){
      while(father[x]!=x)
        x = father[x];
      return x;
    } 

    // Given x and y node
    // Find their root individaully
    // If they're already in the same set, return false 
    // Running time: O(lg n)
    bool merge(int x,int y){
      x = find(x);y = find(y);
      if(x == y)
        return false;
      else{
        if(sz[x] < sz[y]) // weighted merge by size
          std::swap(x,y);
        sz[x] = sz[x]+sz[y];
        father[y] = x;
        return true;
      }
    }
  };
}