UnionFind.h

//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//=========================================================================
 
#ifndef smtk_common_UnionFind_h
#define smtk_common_UnionFind_h
 
#include <map>
#include <set>
#include <vector>
 
namespace smtk
{
namespace common
{
 
template<typename T>
struct UnionFindSet
{
  UnionFindSet(T parent, T rank = 0)
  {
    m_parent = parent;
    m_rank = rank;
  }
  UnionFindSet(const UnionFindSet& other)
  {
    m_parent = other.m_parent;
    m_rank = other.m_rank;
  }
  T m_parent;
  T m_rank;
};
 
template<typename T>
class UnionFind
{
public:
  typedef T value_type;
 
  T newSet();
  T mergeSets(T a, T b);
  T find(T src);
  std::set<T> roots();
  void collapseIds(std::map<T, T>& collapsedIds, T startCount);
  T size() const { return static_cast<T>(m_sets.size()); }
 
  std::vector<UnionFindSet<T>> m_sets;
};
 
template<typename T>
T UnionFind<T>::newSet()
{
  T setId = this->size();
  UnionFindSet<T> entry(setId, 0);
  m_sets.push_back(entry);
  return setId;
}
 
template<typename T>
T UnionFind<T>::mergeSets(T a, T b)
{
  T aRoot = this->find(a);
  T bRoot = this->find(b);
 
  if (aRoot == bRoot)
  {
    return aRoot;
  }
 
  if (aRoot < 0 || bRoot < 0)
  {
    return -1;
  }
 
  T aRank = m_sets[aRoot].m_rank;
  T bRank = m_sets[bRoot].m_rank;
  if (aRank < bRank)
  {
    m_sets[aRoot].m_parent = bRoot;
    return bRoot;
  }
  else if (bRank == aRank)
  {
    ++m_sets[aRoot].m_rank;
  }
  m_sets[bRoot].m_parent = aRoot;
  return aRoot;
}
 
template<typename T>
T UnionFind<T>::find(T src)
{
  if (src < 0 || src >= this->size())
  {
    return -1;
  }
  T parent = m_sets[src].m_parent;
  if (parent != src)
  {
    m_sets[src].m_parent = this->find(parent);
  }
  return m_sets[src].m_parent;
}
 
template<typename T>
std::set<T> UnionFind<T>::roots()
{
  typename std::set<T> roots;
  typename std::vector<UnionFindSet<T>>::iterator it;
  T i = 0;
  for (it = m_sets.begin(); it != m_sets.end(); ++it, ++i)
  {
    if (i == it->m_parent)
    {
      roots.insert(i);
    }
  }
  return roots;
}
 
template<typename T>
void UnionFind<T>::collapseIds(std::map<T, T>& collapsedIds, T startCount)
{
  std::set<T> roots = this->roots();
  typename std::set<T>::iterator it;
  for (it = roots.begin(); it != roots.end(); ++it)
  {
    // Do not relabel any pre-existing entries in collapsedIds.
    typename std::map<T, T>::iterator cit = collapsedIds.find(*it);
    if (cit == collapsedIds.end())
    {
      //cout << "Collapse " << *it << " to " << startCount << "\n";
      collapsedIds[*it] = startCount++;
    }
  }
}
 
} // namespace common
} // namespace smtk
 
#endif // smtk_common_UnionFind_h