// -*- c++ -*- //
#ifndef HUGO_UNION_FIND_H
#define HUGO_UNION_FIND_H

#include <vector>
#include <utility>

namespace hugo {

  template <typename T, typename TIntMap>
  class UnionFind {
    
  public:
    typedef T ElementType;
    typedef std::pair<int,int> PairType;

  private:
    std::vector<PairType> data;
    TIntMap& map;

  public:
    UnionFind(TIntMap& m) : map(m) {}


    int whichComponent(T a)
    {
      int comp0 = map[a];
      if (comp0 < 0) {
	return insertNewElement(a);
      }
      int comp = comp0;
      int next;
      while ( (next = data[comp].first) != comp) {
	comp = next;
      }
      while ( (next = data[comp0].first) != comp) {
	data[comp0].first = comp;
	comp0 = next;
      }

      return comp;
    }

    int insertNewElement(T a)
    {
      int n = data.size();
      data.push_back(make_pair(n, 1));
      map.set(a,n);
      return n;
    }

    bool joinComponents(T a, T b)
    {
      int ca = whichComponent(a);
      int cb = whichComponent(b);

      if ( ca == cb ) 
	return false;

      if ( data[ca].second > data[cb].second ) {
	data[cb].first = ca;
	data[ca].second += data[cb].second;
      }
      else {
	data[ca].first = cb;
	data[cb].second += data[ca].second;
      }
      return true;
    }

    int componentSize(T a)
    {
      int ca = whichComponent(a);
      return data[ca].second;
    }

  };

} //namespace hugo

#endif //HUGO_UNION_FIND_H