Add union-find algorithm to smtk::common.

smtk/common/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 {
+
+/// Internal storage for the UnionFind class.
+template<typename T>
+struct UnionFindSet
+{
+  UnionFindSet(T parent, T rank = 0)
+    {
+    this->m_parent = parent;
+    this->m_rank = rank;
+    }
+  UnionFindSet(const UnionFindSet& other)
+    {
+    this->m_parent = other.m_parent;
+    this->m_rank = other.m_rank;
+    }
+  T m_parent;
+  T m_rank;
+};
+
+/**\brief A disjoint-set structure for fast union-find operations.
+  *
+  * This class is templated on the integer set-id type.
+  * The template parameter should be a signed integer;
+  * smaller sizes may improve performance when a smaller number
+  * of unique sets is acceptable.
+  *
+  * A negative integer is used to identify an invalid value;
+  * for instance, calling Find() on an integer not returned by
+  * NewSet or MergeSets is an error and will return -1.
+  */
+template<typename T>
+class UnionFind
+{
+public:
+  /// Each value used to identify a set is of this type.
+  typedef T value_type;
+
+  /// Add a new set (disjoint from all others).
+  T newSet();
+  /// Connect two disjoint sets, returning the ID of their union (which may not be a or b).
+  T mergeSets(T a, T b);
+  /// Find the parent of a set.
+  T find(T src);
+  /// Get the current root sets (entries in Sets whose parents are themselves)
+  std::set<T> roots();
+  /// Popoulate a map with remaining disjoint sets numbered starting at \a startCount.
+  void collapseIds(std::map<T, T>& collapsedIds, T startCount);
+  /// Return the number of sets that have been created.
+  T size() const
+    { return static_cast<T>(this->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);
+  this->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;
+    }
+
+  T aRank = this->m_sets[aRoot].m_rank;
+  T bRank = this->m_sets[bRoot].m_rank;
+  if (aRank < bRank)
+    {
+    this->m_sets[aRoot].m_parent = bRoot;
+    return bRoot;
+    }
+  else if (bRank == aRank)
+    {
+    ++this->m_sets[aRoot].m_rank;
+    }
+  this->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 = this->m_sets[src].m_parent;
+  if (parent != src)
+    {
+    this->m_sets[src].m_parent = this->find(parent);
+    }
+  return this->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 = this->m_sets.begin(); it != this->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

smtk/common/testing/cxx/CMakeLists.txt

@@ -3,6 +3,7 @@ smtk_public_headers(helpers.h)
 set(commonTests
   unitUUID
   unitPaths
+  unitUnionFind
 )
 
 foreach (test ${commonTests})

smtk/common/testing/cxx/unitUnionFind.cxx

+#include "smtk/common/UnionFind.h"
+
+#include <iostream>
+
+#include "smtk/common/testing/cxx/helpers.h"
+
+using namespace smtk::common;
+
+template<typename T>
+int testUnionFind()
+{
+  UnionFind<T> uf;
+  test(uf.size() == 0, "Expected an initially-empty set.");
+  test(uf.find(0) == -1, "Expected -1 when calling Find() on an invalid set.");
+
+  T s0 = uf.newSet();
+  T s1 = uf.newSet();
+  T s2 = uf.newSet();
+  test(uf.size() == 3, "Expected a set of size 3.");
+
+  test(uf.find(s0) >= 0, "Expected a valid set ID for s0.");
+  test(uf.find(s1) >= 0, "Expected a valid set ID for s1.");
+  test(uf.find(s2) >= 0, "Expected a valid set ID for s2.");
+  test(uf.find(s0) != uf.find(s1), "Expected s0 and s1 to be initially disjoint.");
+  test(uf.find(s0) != uf.find(s2), "Expected s0 and s2 to be initially disjoint.");
+
+  T s3 = uf.mergeSets(s0, s1);
+
+  test(uf.find(s0) == uf.find(s1), "Expected s0 and s1 to match after merge.");
+  test(uf.find(s0) == uf.find(s3), "Expected s0 and s3 to match after merge.");
+  test(uf.find(s1) == uf.find(s3), "Expected s0 and s3 to match after merge.");
+
+  test(uf.find(s0) != uf.find(s2), "Expected s0 and s2 to be disjoint after merge.");
+  test(uf.find(s3) != uf.find(s2), "Expected s3 and s2 to be disjoint after merge.");
+
+  typename std::map<T,T> collapse;
+  uf.collapseIds(collapse, 100);
+
+  std::cout
+    << static_cast<int>(s0) << " -> " << static_cast<int>(collapse[uf.find(s0)]) << "\n"
+    << static_cast<int>(s1) << " -> " << static_cast<int>(collapse[uf.find(s1)]) << "\n"
+    << static_cast<int>(s2) << " -> " << static_cast<int>(collapse[uf.find(s2)]) << "\n"
+    << "\n"
+    ;
+
+  test(collapse[uf.find(s0)] == 100, "Invalid collapsed ID for s0.");
+  test(collapse[uf.find(s1)] == 100, "Invalid collapsed ID for s1.");
+  test(collapse[uf.find(s2)] == 101, "Invalid collapsed ID for s2.");
+  test(collapse.size() == uf.roots().size(), "Expected a collapsed-ID map of size 2.");
+  test(uf.size() == 3, "Still expected a set of size 3.");
+
+  return 0;
+}
+
+int main(int argc, char* argv[])
+{
+  return
+    testUnionFind<int>() ||
+    testUnionFind<signed char>();
+}