hphp/util/union-find.h

   1 /*
   2
   3   +----------------------------------------------------------------------+
   4   | HipHop for PHP                                                       |
   5   +----------------------------------------------------------------------+
   6   | Copyright (c) 2010-present Facebook, Inc. (http://www.facebook.com)  |
   7   +----------------------------------------------------------------------+
   8   | This source file is subject to version 3.01 of the PHP license,      |
   9   | that is bundled with this package in the file LICENSE, and is        |
  10   | available through the world-wide-web at the following url:           |
  11   | http://www.php.net/license/3_01.txt                                  |
  12   | If you did not receive a copy of the PHP license and are unable to   |
  13   | obtain it through the world-wide-web, please send a note to          |
  14   | license@php.net so we can mail you a copy immediately.               |
  15   +----------------------------------------------------------------------+
  16 */
  17
  18 #pragma once
  19
  20 #include <vector>
  21 #include <folly/container/F14Map.h>
  22
  23 #include "hphp/util/assertions.h"
  24
  25 namespace HPHP {
  26
  27 template <typename T>
  28 struct UnionFind {
  29   // Create a disjoint set containing only the given node. Return its index.
  30   size_t insert(const T& node) {
  31     auto const index = nodes.size();
  32     nodes.push_back({index, 1});
  33     auto const pair = indices.emplace(node, index);
  34     assertx(pair.second);
  35     return pair.first->second;
  36   }
  37
  38   // Return the index of the given node, assigned sequentially at insertion.
  39   folly::Optional<size_t> lookup(const T& node) const {
  40     auto const it = indices.find(node);
  41     return it == indices.end() ? folly::none : folly::make_optional(it->second);
  42   }
  43
  44   // Merge the disjoint sets containing the two nodes. The nodes must have
  45   // been inserted into the data structure already.
  46   void merge(const T& node1, const T& node2) {
  47     assertx(indices.contains(node1));
  48     assertx(indices.contains(node2));
  49     if (node1 == node2) return;
  50     auto index1 = find(indices.at(node1));
  51     auto index2 = find(indices.at(node2));
  52     if (index1 == index2) return;
  53
  54     if (nodes[index1].size < nodes[index2].size) {
  55       std::swap(index1, index2);
  56     }
  57     nodes[index2].parent = index1;
  58     nodes[index1].size += nodes[index2].size;
  59   }
  60
  61   // Return the number of disjoint sets.
  62   size_t countGroups() const {
  63     auto result = size_t{0};
  64     for (auto i = 0; i < nodes.size(); i++) {
  65       if (nodes[i].parent == i) result++;
  66     }
  67     return result;
  68   }
  69
  70   // Execute `f` for each disjoint set. `f` takes std::vector<T>&.
  71   template <typename F>
  72   void forEachGroup(F&& f) {
  73     DEBUG_ONLY auto const count = countGroups();
  74     folly::F14FastMap<size_t, std::vector<T>> groups;
  75     for (auto const& pair : indices) {
  76       groups[find(pair.second)].emplace_back(pair.first);
  77     }
  78     assertx(count == countGroups());
  79     assertx(count == groups.size());
  80     for (auto& pair : groups) {
  81       f(pair.second);
  82     }
  83   }
  84
  85 private:
  86   size_t find(size_t index) {
  87     while (nodes[index].parent != index) {
  88       auto const grandparent = nodes[nodes[index].parent].parent;
  89       index = nodes[index].parent = grandparent;
  90     }
  91     return index;
  92   }
  93
  94   struct Node {
  95     size_t parent;
  96     size_t size;
  97   };
  98   std::vector<Node> nodes;
  99   folly::F14FastMap<T, size_t> indices;
 100 };
 101
 102 }