beckerjc@150: // -*- c++ -*- //
beckerjc@218: #ifndef HUGO_UNION_FIND_H
beckerjc@218: #define HUGO_UNION_FIND_H
beckerjc@150: 
beckerjc@150: #include <vector>
beckerjc@150: #include <utility>
beckerjc@150: 
beckerjc@150: namespace hugo {
beckerjc@150: 
beckerjc@150:   template <typename T, typename TIntMap>
beckerjc@150:   class UnionFind {
beckerjc@150:     
beckerjc@150:   public:
beckerjc@150:     typedef T ElementType;
beckerjc@150:     typedef std::pair<int,int> PairType;
beckerjc@150: 
beckerjc@150:   private:
beckerjc@150:     std::vector<PairType> data;
beckerjc@150:     TIntMap& map;
beckerjc@150: 
beckerjc@150:   public:
beckerjc@150:     UnionFind(TIntMap& m) : map(m) {}
beckerjc@150: 
beckerjc@150: 
beckerjc@150:     int whichComponent(T a)
beckerjc@150:     {
beckerjc@150:       int comp0 = map.get(a);
beckerjc@150:       if (comp0 < 0) {
beckerjc@150: 	return insertNewElement(a);
beckerjc@150:       }
beckerjc@150:       int comp = comp0;
beckerjc@150:       int next;
beckerjc@150:       while ( (next = data[comp].first) != comp) {
beckerjc@150: 	comp = next;
beckerjc@150:       }
beckerjc@150:       while ( (next = data[comp0].first) != comp) {
beckerjc@150: 	data[comp0].first = comp;
beckerjc@150: 	comp0 = next;
beckerjc@150:       }
beckerjc@150: 
beckerjc@150:       return comp;
beckerjc@150:     }
beckerjc@150: 
beckerjc@150:     int insertNewElement(T a)
beckerjc@150:     {
beckerjc@150:       int n = data.size();
beckerjc@150:       data.push_back(make_pair(n, 1));
beckerjc@150:       map.set(a,n);
beckerjc@150:       return n;
beckerjc@150:     }
beckerjc@150: 
beckerjc@150:     bool joinComponents(T a, T b)
beckerjc@150:     {
beckerjc@150:       int ca = whichComponent(a);
beckerjc@150:       int cb = whichComponent(b);
beckerjc@150: 
beckerjc@150:       if ( ca == cb ) 
beckerjc@150: 	return false;
beckerjc@150: 
beckerjc@150:       if ( data[ca].second > data[cb].second ) {
beckerjc@150: 	data[cb].first = ca;
beckerjc@150: 	data[ca].second += data[cb].second;
beckerjc@150:       }
beckerjc@150:       else {
beckerjc@150: 	data[ca].first = cb;
beckerjc@150: 	data[cb].second += data[ca].second;
beckerjc@150:       }
beckerjc@150:       return true;
beckerjc@150:     }
beckerjc@150: 
beckerjc@218:     int componentSize(T a)
beckerjc@218:     {
beckerjc@218:       int ca = whichComponent(a);
beckerjc@218:       return data[ca].second;
beckerjc@218:     }
beckerjc@218: 
beckerjc@150:   };
beckerjc@150: 
beckerjc@150: } //namespace hugo
beckerjc@150: 
beckerjc@218: #endif //HUGO_UNION_FIND_H