dsu.h

Go to the documentation of this file.

#pragma once
 
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <vector>
 
namespace clustering {

template <class Idx = std::uint32_t> class UnionFind {
public:
  static_assert(std::is_unsigned_v<Idx>, "UnionFind Idx must be an unsigned integer type");

  explicit UnionFind(std::size_t n) : m_parent(n), m_rank(n, 0), m_size(n, 1), m_components(n) {
    for (std::size_t i = 0; i < n; ++i) {
      m_parent[i] = static_cast<Idx>(i);
    }
  }

  Idx find(Idx x) noexcept {
    assert(static_cast<std::size_t>(x) < m_parent.size() && "UnionFind::find index out of range");
    Idx root = x;
    while (m_parent[root] != root) {
      root = m_parent[root];
    }
    Idx node = x;
    while (m_parent[node] != root) {
      const Idx next = m_parent[node];
      m_parent[node] = root;
      node = next;
    }
    return root;
  }

  bool unite(Idx a, Idx b) noexcept {
    Idx ra = find(a);
    Idx rb = find(b);
    if (ra == rb) {
      return false;
    }
    if (m_rank[ra] < m_rank[rb]) {
      m_parent[ra] = rb;
      m_size[rb] += m_size[ra];
    } else if (m_rank[ra] > m_rank[rb]) {
      m_parent[rb] = ra;
      m_size[ra] += m_size[rb];
    } else {
      m_parent[rb] = ra;
      m_size[ra] += m_size[rb];
      ++m_rank[ra];
    }
    --m_components;
    return true;
  }

  bool sameComponent(Idx a, Idx b) noexcept { return find(a) == find(b); }

  [[nodiscard]] std::size_t countComponents() const noexcept { return m_components; }

  [[nodiscard]] std::size_t size() const noexcept { return m_parent.size(); }

  [[nodiscard]] std::size_t componentSize(Idx root) const noexcept {
    assert(static_cast<std::size_t>(root) < m_parent.size() &&
           "UnionFind::componentSize index out of range");
    return m_size[root];
  }
 
private:
  std::vector<Idx> m_parent;
  std::vector<std::uint8_t> m_rank;
  std::vector<std::size_t> m_size;
  std::size_t m_components;
};
 
} // namespace clustering