yusuke
1/12/2020 - 7:40 AM

Union Find Tree

struct UnionFindTree {
  std::vector<int> data;

  UnionFindTree(int size) : data(size, -1) { }
  bool Unite(int x, int y) {
    x = Root(x); y = Root(y);
    if (x != y) {
      if (data[y] < data[x]) std::swap(x, y);
      data[x] += data[y];
      data[y] = x;
    }
    return x != y;
  }
  bool IsSameRoot(int x, int y) {
    return Root(x) == Root(y);
  }
  int Root(int x) {
    return data[x] < 0 ? x : data[x] = Root(data[x]);
  }
  int Size(int x) { return -data[Root(x)]; }
};

How to use:

int main() {
  int N; cin >> N;
  UnionFindTree tree(N);
  
  // Build Union find tree.
  int E; cin >> E;
  for (int e = 0; e < E; ++e) {
    int a, b; cin >> a >> b;
    tree.Unite(a, b);
  }
  
  // Usage1: check node#3 and node#5 is connected or not.
  if (tree.IsSameRoot(3, 5)) cout << "Node 3 and 5 are connected."
  else                       cout << "Node 3 and 5 are NOT conenected"
  
  // Usage2: calculate node_num connected to node#1.
  cout << "Num of node connected to node#1: " << tree.Size(1) << endl;
}