A quick implementation of disjoint sets on a dense set of indices. O(1) `merge`, and after `reduce_indices` O(1) `is_same_set`. Even `reduce_indices` is O(N)

IndexDisjointSet.hpp

#pragma once
#include <ranges>
#include <vector>

/// IndexDisjointSet lets you merge indices into disjoint subsets
// via merge and then query which indices belong in the same set 
// with is_same_set.
// Merges are always O(1), and calling reduce_indices sets 
// is_same_set to be O(1).
// Minor tweaks to merge or is_same_set so they reduce the depth
// makes everything is amortized O(1).
class IndexDisjointSet {
public:
  // initialize all elements as [0,1,2,3...]
  IndexDisjointSet(size_t size)
      : m_list(std::views::iota(size_t{}, size_t{size}) |
               std::ranges::to<std::vector<size_t>>()) {}

  /// main query of whether two indices are in the same set
  [[nodiscard]] auto is_same_set(size_t a, size_t b) const {
    return get_root(a) == get_root(b);
  }

  /// Finds canonical root id for each node
  [[nodiscard]] auto get_root(std::size_t x) const -> std::size_t {
    while (m_list[x] != x) {
      x = m_list[x];
    }
    return x;
  }

  /// Compactifies indices so m_list values are the new compact indices
  /// After this is_same_set is O(1)
  void reduce_indices() {
    for (std::size_t &k : m_list) {
      k = m_list[k];
    }
  }

  /// Merges two sets to one set
  auto merge(size_t a, size_t b) -> size_t {
    // Find the roots of both nodes, makes the larger root point to the smaller
    // one
    a = get_root(a);
    b = get_root(b);
    size_t root;
    if (a < b) {
      root = a;
      m_list[b] = a;
    } else {
      root = b;
      m_list[a] = b;
    }
    return root;
  }

private:
  std::vector<size_t> m_list;
};