[base] ClusteringMap.

base/CMakeLists.txt

@@ -13,6 +13,7 @@ set(
   cache.hpp
   cancellable.hpp
   checked_cast.hpp
+  clustering_map.hpp
   collection_cast.hpp
   condition.cpp
   condition.hpp

base/base.pro

@@ -49,6 +49,7 @@ HEADERS += \
     cache.hpp \
     cancellable.hpp \
     checked_cast.hpp \
+    clustering_map.hpp \
     collection_cast.hpp \
     condition.hpp \
     deferred_task.hpp \

base/base_tests/CMakeLists.txt

@@ -9,6 +9,7 @@ set(
   buffer_vector_test.cpp
   bwt_tests.cpp
   cache_test.cpp
+  clustering_map_tests.cpp
   collection_cast_test.cpp
   condition_test.cpp
   containers_test.cpp

base/base_tests/base_tests.pro

@@ -19,6 +19,7 @@ SOURCES += \
   buffer_vector_test.cpp \
   bwt_tests.cpp \
   cache_test.cpp \
+  clustering_map_tests.cpp \
   collection_cast_test.cpp \
   condition_test.cpp \
   containers_test.cpp \

base/base_tests/clustering_map_tests.cpp

@@ -0,0 +1,84 @@
+#include "testing/testing.hpp"
+
+#include "base/clustering_map.hpp"
+
+#include <algorithm>
+#include <string>
+#include <utility>
+#include <vector>
+
+using namespace base;
+using namespace std;
+
+namespace
+{
+template <typename T>
+vector<T> Sort(vector<T> vs)
+{
+  sort(vs.begin(), vs.end());
+  return vs;
+}
+
+template <typename Key, typename Value, typename Hash = std::hash<Key>>
+class ClusteringMapAdapter
+{
+public:
+  template <typename V>
+  void Append(Key const & key, V && value)
+  {
+    m_m.Append(key, std::forward<V>(value));
+  }
+
+  void Union(Key const & u, Key const & v) { m_m.Union(u, v); }
+
+  std::vector<Value> Get(Key const & key) { return Sort(m_m.Get(key)); }
+
+private:
+  ClusteringMap<Key, Value, Hash> m_m;
+};
+
+UNIT_TEST(ClusteringMap_Smoke)
+{
+  {
+    ClusteringMapAdapter<int, string> m;
+    TEST(m.Get(0).empty(), ());
+    TEST(m.Get(1).empty(), ());
+
+    m.Union(0, 1);
+    TEST(m.Get(0).empty(), ());
+    TEST(m.Get(1).empty(), ());
+  }
+
+  {
+    ClusteringMapAdapter<int, string> m;
+    m.Append(0, "Hello");
+    m.Append(1, "World!");
+
+    TEST_EQUAL(m.Get(0), vector<string>({"Hello"}), ());
+    TEST_EQUAL(m.Get(1), vector<string>({"World!"}), ());
+
+    m.Union(0, 1);
+    TEST_EQUAL(m.Get(0), vector<string>({"Hello", "World!"}), ());
+    TEST_EQUAL(m.Get(1), vector<string>({"Hello", "World!"}), ());
+
+    m.Append(2, "alpha");
+    m.Append(3, "beta");
+    m.Append(4, "gamma");
+
+    TEST_EQUAL(m.Get(2), vector<string>({"alpha"}), ());
+    TEST_EQUAL(m.Get(3), vector<string>({"beta"}), ());
+    TEST_EQUAL(m.Get(4), vector<string>({"gamma"}), ());
+
+    m.Union(2, 3);
+    m.Union(3, 4);
+
+    TEST_EQUAL(m.Get(2), vector<string>({"alpha", "beta", "gamma"}), ());
+    TEST_EQUAL(m.Get(3), vector<string>({"alpha", "beta", "gamma"}), ());
+    TEST_EQUAL(m.Get(4), vector<string>({"alpha", "beta", "gamma"}), ());
+
+    TEST_EQUAL(m.Get(5), vector<string>(), ());
+    m.Union(2, 5);
+    TEST_EQUAL(m.Get(5), vector<string>({"alpha", "beta", "gamma"}), ());
+  }
+}
+}  // namespace

base/clustering_map.hpp

@@ -0,0 +1,107 @@
+#pragma once
+
+#include "base/assert.hpp"
+
+#include <cstddef>
+#include <functional>
+#include <utility>
+#include <vector>
+
+namespace base
+{
+// Maps keys to lists of values, but allows to clusterize keys
+// together, and to get all values from a cluster.
+//
+// NOTE: the map is NOT thread-safe.
+template <typename Key, typename Value, typename Hash = std::hash<Key>>
+class ClusteringMap
+{
+public:
+  // Appends |value| to the list of values in the cluster
+  // corresponding to |key|.
+  //
+  // Amortized complexity: O(log*(n) * F), where n is the total number
+  // of keys in the map, F is the complexity of find in unordered_map.
+  template <typename V>
+  void Append(Key const & key, V && value)
+  {
+    auto & entry = GetRoot(key);
+    entry.m_values.push_back(std::forward<V>(value));
+  }
+
+  // Unions clusters corresponding to |u| and |v|.
+  //
+  // Amortized complexity: O(log*(n) * F + log(m)), where n is the
+  // total number of keys and m is the total number of values in the
+  // map, F is the complexity of find in unordered_map.
+  void Union(Key const & u, Key const & v)
+  {
+    auto & ru = GetRoot(u);
+    auto & rv = GetRoot(v);
+    if (ru.m_root == rv.m_root)
+      return;
+
+    if (ru.m_rank < rv.m_rank)
+      Attach(rv /* root */, ru /* child */);
+    else
+      Attach(ru /* root */, rv /* child */);
+  }
+
+  // Returns all values from the cluster corresponding to |key|.
+  //
+  // Amortized complexity: O(log*(n) * F), where n is the total number
+  // of keys in the map, F is the complexity of find in unordered map.
+  std::vector<Value> const & Get(Key const & key)
+  {
+    auto const & entry = GetRoot(key);
+    return entry.m_values;
+  }
+
+private:
+  struct Entry
+  {
+    Key m_root;
+    size_t m_rank = 0;
+    std::vector<Value> m_values;
+  };
+
+  Entry & GetRoot(Key const & key)
+  {
+    auto & entry = GetEntry(key);
+    if (entry.m_root == key)
+      return entry;
+
+    auto & root = GetRoot(entry.m_root);
+    entry.m_root = root.m_root;
+    return root;
+  }
+
+  void Attach(Entry & parent, Entry & child)
+  {
+    ASSERT_LESS_OR_EQUAL(child.m_rank, parent.m_rank, ());
+
+    child.m_root = parent.m_root;
+    if (child.m_rank == parent.m_rank)
+      ++parent.m_rank;
+
+    auto & pv = parent.m_values;
+    auto & cv = child.m_values;
+    if (pv.size() < cv.size())
+      pv.swap(cv);
+    pv.insert(pv.end(), cv.begin(), cv.end());
+  }
+
+  Entry & GetEntry(Key const & key)
+  {
+    auto it = m_table.find(key);
+    if (it != m_table.end())
+      return it->second;
+
+    auto & entry = m_table[key];
+    entry.m_root = key;
+    return entry;
+  }
+
+  std::unordered_map<Key, Entry, Hash> m_table;
+};
+}  // namespace base