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 #include "hphp/util/optional.h"
  25
  26 namespace HPHP {
  27
  28 template <typename T>
  29 struct UnionFind {
  30   // Create a disjoint set containing only the given node. Return its index.
  31   size_t insert(const T& node) {
  32     auto const index = nodes.size();
  33     nodes.push_back({index, 1});
  34     auto const pair = indices.emplace(node, index);
  35     assertx(pair.second);
  36     return pair.first->second;
  37   }
  38
  39   // Return the index of the given node, assigned sequentially at insertion.
  40   Optional<size_t> lookup(const T& node) const {
  41     auto const it = indices.find(node);
  42     return it == indices.end() ? std::nullopt : make_optional(it->second);
  43   }
  44
  45   // Merge the disjoint sets containing the two nodes. The nodes must have
  46   // been inserted into the data structure already.
  47   void merge(const T& node1, const T& node2) {
  48     assertx(indices.contains(node1));
  49     assertx(indices.contains(node2));
  50     if (node1 == node2) return;
  51     auto index1 = find(indices.at(node1));
  52     auto index2 = find(indices.at(node2));
  53     if (index1 == index2) return;
  54
  55     if (nodes[index1].size < nodes[index2].size) {
  56       std::swap(index1, index2);
  57     }
  58     nodes[index2].parent = index1;
  59     nodes[index1].size += nodes[index2].size;
  60   }
  61
  62   // Return the number of disjoint sets.
  63   size_t countGroups() const {
  64     auto result = size_t{0};
  65     for (auto i = 0; i < nodes.size(); i++) {
  66       if (nodes[i].parent == i) result++;
  67     }
  68     return result;
  69   }
  70
  71   // Execute `f` for each disjoint set. `f` takes std::vector<T>&.
  72   template <typename F>
  73   void forEachGroup(F&& f) {
  74     DEBUG_ONLY auto const count = countGroups();
  75     folly::F14FastMap<size_t, std::vector<T>> groups;
  76     for (auto const& pair : indices) {
  77       groups[find(pair.second)].emplace_back(pair.first);
  78     }
  79     assertx(count == countGroups());
  80     assertx(count == groups.size());
  81     for (auto& pair : groups) {
  82       f(pair.second);
  83     }
  84   }
  85
  86 private:
  87   size_t find(size_t index) {
  88     while (nodes[index].parent != index) {
  89       auto const grandparent = nodes[nodes[index].parent].parent;
  90       index = nodes[index].parent = grandparent;
  91     }
  92     return index;
  93   }
  94
  95   struct Node {
  96     size_t parent;
  97     size_t size;
  98   };
  99   std::vector<Node> nodes;
 100   folly::F14FastMap<T, size_t> indices;
 101 };
 102
 103 }