dsu.h

Go to the documentation of this file.

#pragma once
 
#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#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;
};

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

  explicit AtomicUnionFind(std::size_t n) : m_parent(n) {
    for (std::size_t i = 0; i < n; ++i) {
      m_parent[i].store(static_cast<Idx>(i), std::memory_order_relaxed);
    }
  }

  Idx find(Idx x) noexcept {
    assert(static_cast<std::size_t>(x) < m_parent.size() &&
           "AtomicUnionFind::find index out of range");
    while (true) {
      Idx parent = m_parent[x].load(std::memory_order_acquire);
      if (parent == x) {
        return x;
      }
      const Idx grandparent = m_parent[parent].load(std::memory_order_acquire);
      if (grandparent == parent) {
        return parent;
      }
      // Halve the path: retarget x at its grandparent. A lost race means another thread
      // already advanced the slot; the walk continues from the grandparent either way.
      m_parent[x].compare_exchange_weak(parent, grandparent, std::memory_order_release,
                                        std::memory_order_relaxed);
      x = grandparent;
    }
  }

  bool unite(Idx a, Idx b) noexcept {
    while (true) {
      Idx ra = find(a);
      Idx rb = find(b);
      if (ra == rb) {
        return false;
      }
      if (rb < ra) {
        std::swap(ra, rb);
      }
      Idx expected = rb;
      if (m_parent[rb].compare_exchange_strong(expected, ra, std::memory_order_acq_rel,
                                               std::memory_order_acquire)) {
        return true;
      }
      // rb stopped being a root mid-flight; restart from the advanced positions.
      a = ra;
      b = expected;
    }
  }

  [[nodiscard]] std::size_t size() const noexcept { return m_parent.size(); }
 
private:
  std::vector<std::atomic<Idx>> m_parent;
};
 
} // namespace clustering