boruvka_mst_backend.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "clustering/always_assert.h"
#include "clustering/hdbscan/detail/boruvka_traversal.h"
#include "clustering/hdbscan/mst_output.h"
#include "clustering/index/kdtree.h"
#include "clustering/math/dsu.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering::hdbscan {

template <class T> class BoruvkaMstBackend {
  static_assert(
      std::is_same_v<T, float>,
      "BoruvkaMstBackend<T> supports only float; a double specialization is out of scope.");
 
public:
  BoruvkaMstBackend() = default;

  void run(const NDArray<T, 2> &X, std::size_t minSamples, math::Pool pool, MstOutput<T> &out) {
    const std::size_t n = X.dim(0);
    CLUSTERING_ALWAYS_ASSERT(minSamples >= 1);
    CLUSTERING_ALWAYS_ASSERT(minSamples < n);
 
    out.edges.clear();
    out.edges.reserve(n - 1);
    out.coreDistances = NDArray<T, 1>(std::array<std::size_t, 1>{n});
 
    // Build the tree once; reused for the kNN core-distance query and for every Boruvka
    // round's per-point nearest-out-of-component traversal. LeafSize=32 halves the internal
    // node count against the default 16-leaf tree; at d>=4 the batched-distance leaf kernel
    // amortises the larger leaf's arithmetic well enough that the saved per-node AABB gap and
    // pivot distance wins net positive on every shape exercised by HDBSCAN's dispatch gates.
    const KDTree<T, KDTreeDistanceType::kEucledian, 64> tree(X);
 
    // Phase 1: core distances via kNN at k = minSamples. knnQuery returns neighbours sorted
    // ascending by squared distance, so the minSamples-th (1-based) neighbour is at column
    // @c minSamples - 1. Squared distances are the in-memory unit across the pipeline; the
    // downstream MRD weight and Prim oracle use the same scale.
    const auto kSigned = static_cast<std::int32_t>(minSamples);
    auto [knnIdx, knnSqDist] = tree.knnQuery(kSigned, pool);
    T *coreDistData = out.coreDistances.data();
    for (std::size_t i = 0; i < n; ++i) {
      coreDistData[i] = knnSqDist(i, minSamples - 1);
    }
 
    // Phase 2: Boruvka rounds. Each round computes one best out-of-component candidate edge
    // per component via the per-point traversal, then unions those edges in ascending order of
    // weight. Processing in ascending weight keeps the MRD-MST total identical to the classical
    // Kruskal/Prim result, modulo equal-weight ties which admit multiple valid MSTs.
    UnionFind<std::uint32_t> uf(n);
    std::vector<std::int32_t> componentOf(n, std::int32_t{0});
    std::vector<detail::ComponentBestEdge<T>> bestPerComponent(n);
    std::vector<std::int32_t> activeRoots;
    activeRoots.reserve(n);
 
    bool firstRound = true;
    while (uf.countComponents() > 1) {
      for (std::size_t i = 0; i < n; ++i) {
        componentOf[i] = static_cast<std::int32_t>(uf.find(static_cast<std::uint32_t>(i)));
      }
 
      // First-round fast path: every point is its own component; the kNN seed inside
      // @c nearestOutComponent picks the right answer without a tree walk.
      detail::nearestOutComponent<T>(tree, std::span<const std::int32_t>(componentOf),
                                     out.coreDistances, knnIdx, knnSqDist, pool, bestPerComponent,
                                     /*allSingletonComponents=*/firstRound);
      firstRound = false;
 
      // Collect unique root candidates that produced a finite best. A component without a
      // finite best on a >1-component graph would be a correctness failure -- the fan-out
      // must have reached every component's points. Assert rather than silently skip so the
      // bug surfaces at the failing round.
      activeRoots.clear();
      for (std::size_t i = 0; i < n; ++i) {
        if (std::cmp_equal(componentOf[i], i)) {
          activeRoots.push_back(static_cast<std::int32_t>(i));
        }
      }
 
      // Apply in ascending-weight order. Tie-break by (u, v) keeps the output deterministic
      // across equal-weight candidates that may arrive from different components.
      std::sort(activeRoots.begin(), activeRoots.end(),
                [&](std::int32_t a, std::int32_t b) noexcept {
                  const auto &ea = bestPerComponent[static_cast<std::size_t>(a)];
                  const auto &eb = bestPerComponent[static_cast<std::size_t>(b)];
                  if (ea.weight != eb.weight) {
                    return ea.weight < eb.weight;
                  }
                  if (ea.u != eb.u) {
                    return ea.u < eb.u;
                  }
                  return ea.v < eb.v;
                });
 
      bool unitedAny = false;
      for (const std::int32_t c : activeRoots) {
        const auto &edge = bestPerComponent[static_cast<std::size_t>(c)];
        CLUSTERING_ALWAYS_ASSERT(edge.u >= 0 && edge.v >= 0);
        if (uf.unite(static_cast<std::uint32_t>(edge.u), static_cast<std::uint32_t>(edge.v))) {
          out.edges.push_back(MstEdge<T>{edge.u, edge.v, edge.weight});
          unitedAny = true;
        }
      }
 
      // Safety gate: a round that fails to unite any component would loop forever. This cannot
      // happen on a connected graph when every component has a finite best, but assert rather
      // than infinite-loop if the invariant ever breaks.
      CLUSTERING_ALWAYS_ASSERT(unitedAny);
    }
 
    CLUSTERING_ALWAYS_ASSERT(out.edges.size() == n - 1);
 
    // Canonical output ordering: ascending by weight. The single-linkage tree construction
    // downstream consumes edges in this order, matching the Prim backend's implicit Prim-pop
    // ordering (which is weight-monotone outside of tie groups). Sort explicitly here so the
    // output contract is uniform across backends.
    std::sort(out.edges.begin(), out.edges.end(),
              [](const MstEdge<T> &a, const MstEdge<T> &b) noexcept {
                if (a.weight != b.weight) {
                  return a.weight < b.weight;
                }
                if (a.u != b.u) {
                  return a.u < b.u;
                }
                return a.v < b.v;
              });
  }
};
 
} // namespace clustering::hdbscan