kdtree.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <numeric>
#include <span>
#include <utility>
#include <vector>
 
#include "clustering/always_assert.h"
#include "clustering/index/range_query.h"
#include "clustering/math/aabb.h"
#include "clustering/math/detail/avx2_helpers.h"
#include "clustering/math/detail/radius_scan.h"
#include "clustering/math/detail/sq_distances_block.h"
#include "clustering/math/detail/top_k_neighbors.h"
#include "clustering/math/thread.h"
#include "clustering/memory/linear_alloc.h"
#include "clustering/ndarray.h"
 
namespace clustering {

struct KDTreeNode {
  std::size_t m_index;
  std::size_t m_dim;
  KDTreeNode *m_left;
  KDTreeNode *m_right;
  std::uint32_t m_id;
};

enum class KDTreeDistanceType : std::uint8_t {
  kEucledian 
};

template <class T, KDTreeDistanceType distanceType = KDTreeDistanceType::kEucledian,
          std::size_t LeafSize = 16, class AllocT = LinearAllocator<KDTreeNode>>
class KDTree {
public:
  using value_type = T; 

  KDTree(const NDArray<T, 2> &points, math::Pool pool = {})
      : m_allocator(calculatePoolSize(points.dim(0))), m_points(points), m_dim(points.dim(1)) {
    CLUSTERING_ALWAYS_ASSERT(points.isContiguous());
    const std::size_t n = points.dim(0);
    m_indices.resize(n);
    std::iota(m_indices.begin(), m_indices.end(), 0);
    // Subtree node counts are a pure function of range size, so every node's arena slot and
    // id are known before recursion starts; parallel subtree builds write disjoint slots with
    // no shared allocation state and reproduce the serial layout exactly.
    const std::size_t totalNodes = nodeCountFor(n);
    KDTreeNode *arena = (totalNodes > 0) ? m_allocator.allocate(totalNodes) : nullptr;
    m_root = (totalNodes > 0) ? buildAt(0, n, 0, arena, 0, pool) : nullptr;
    m_nextNodeId = static_cast<std::uint32_t>(totalNodes);
    // Materialize points in tree-build order. After @c buildAt rewrites @c m_indices into a
    // permutation matching the tree layout, a leaf's points live at @c m_points_reordered slots
    // `[leaf.m_index, leaf.m_index + leaf.m_dim)`. Contiguous access there replaces the
    // scatter-indirection `m_points[m_indices[k]`] had, which at low d was the cache-miss
    // ceiling: every leaf-brute-force iteration landed on a random row of @c m_points.
    // Rows land in disjoint slots, so the copy fans out over the pool.
    m_points_reordered.resize(n * m_dim);
    const T *src = points.data();
    T *dst = m_points_reordered.data();
    pool.parallelForBlocks(std::size_t{0}, n, std::size_t{0}, [&](std::size_t lo, std::size_t hi) {
      for (std::size_t k = lo; k < hi; ++k) {
        const T *s = src + (m_indices[k] * m_dim);
        T *d = dst + (k * m_dim);
        for (std::size_t j = 0; j < m_dim; ++j) {
          d[j] = s[j];
        }
      }
    });
    // Populate per-node axis-aligned bounding boxes once the tree is built and the reordered
    // points buffer is materialized. Layout is `(numNodes * 2 * d)` flat: row @c 2*id holds the
    // min-coords vector, row @c 2*id + 1 holds the max-coords vector. Dual-tree walkers consume
    // this through @ref nodeBounds as a pair of @c std::span views; leaving @ref KDTreeNode's
    // size unchanged past the monotonic @c m_id keeps leaf-scan cache behaviour stable at
    // high @c d.
    //
    // Leaf boxes carry the point sweep, so they fan out over the arena; internal boxes are
    // unions of their children, and the post-order ids let a flat ascending-id pass visit
    // every child before its parent without recursion.
    m_nodeBounds.assign(static_cast<std::size_t>(m_nextNodeId) * 2 * m_dim, T{});
    if (m_root != nullptr) {
      KDTreeNode *arenaNodes = m_root;
      pool.parallelForBlocks(std::size_t{0}, totalNodes, std::size_t{0},
                             [&](std::size_t lo, std::size_t hi) {
                               for (std::size_t s = lo; s < hi; ++s) {
                                 const KDTreeNode &node = arenaNodes[s];
                                 if (node.m_left == nullptr && node.m_right == nullptr) {
                                   populateLeafBounds(node);
                                 }
                               }
                             });
      std::vector<const KDTreeNode *> byId(totalNodes, nullptr);
      for (std::size_t s = 0; s < totalNodes; ++s) {
        byId[arenaNodes[s].m_id] = &arenaNodes[s];
      }
      for (std::size_t id = 0; id < totalNodes; ++id) {
        const KDTreeNode &node = *byId[id];
        if (node.m_left != nullptr || node.m_right != nullptr) {
          unionChildBounds(node);
        }
      }
    }
  }

  std::vector<std::size_t> query(const NDArray<T, 1> &query_point, T radius,
                                 std::int64_t limit = -1) const {
    std::vector<std::size_t> indices;
    const T radius_sq = radius * radius;
    ensureLeafSoa();
    std::vector<KDTreeNode *> stack;
    stack.reserve(kDefaultStackReserve);
    // Copy the borrowed row into a scratch buffer so the core walker can assume a contiguous
    // `d`-element block regardless of the source layout. The copy is d stores -- free at d=2.
    std::vector<T> qbuf(m_dim);
    for (std::size_t k = 0; k < m_dim; ++k) {
      qbuf[k] = query_point[k];
    }
    queryImpl(m_root, qbuf.data(), radius_sq, indices, stack, limit);
    return indices;
  }

  [[nodiscard]] index::CoreAdjacency query(T radius, std::size_t minPts, math::Pool pool) const {
    const std::size_t n = m_points.dim(0);
    index::CoreAdjacency out;
    out.rows.resize(n);
    out.isCore.assign(n, 0);
    if (n == 0) {
      return out;
    }
 
    const T radius_sq = radius * radius;
    ensureLeafSoa(pool);
 
    // Above the dim floor, one box-pruned walk per leaf replaces one walk per point: the
    // visited-node count collapses by the leaf occupancy while the pair tests grow only by the
    // box inflation, which the higher per-pair cost at larger d amortizes. Below the floor the
    // per-pair scan is a few ops and the inflation dominates, so points walk individually.
    if (m_dim >= kBlockQueryDimFloor && m_root != nullptr) {
      blockQuery(radius_sq, minPts, pool, out);
      return out;
    }
 
    // Leaves whose bounding-box diagonal fits inside the radius are cliques: every member is
    // within eps of every other, and with at least minPts members every member is a core. A
    // query ball that swallows such a leaf can take one representative edge instead of
    // materializing the members, so flag them once up front.
    const std::size_t totalNodes = nodeCount();
    const KDTreeNode *arenaNodes = m_root;
    std::vector<std::uint8_t> leafAllCore(totalNodes, 0);
    pool.parallelForBlocks(
        std::size_t{0}, totalNodes, std::size_t{0}, [&](std::size_t lo, std::size_t hi) {
          for (std::size_t nodeSlot = lo; nodeSlot < hi; ++nodeSlot) {
            const KDTreeNode &node = arenaNodes[nodeSlot];
            if (node.m_left != nullptr || node.m_right != nullptr || node.m_dim < minPts) {
              continue;
            }
            const auto [bmin, bmax] = nodeBounds(&node);
            T diagSq = T{0};
            for (std::size_t j = 0; j < m_dim; ++j) {
              const T ext = bmax[j] - bmin[j];
              diagSq += ext * ext;
            }
            if (diagSq <= radius_sq) {
              leafAllCore[node.m_id] = 1;
            }
          }
        });
 
    const std::size_t workers = pool.workerCount();
    std::vector<std::vector<std::pair<std::int32_t, std::int32_t>>> workerEdges(workers);
 
    auto runRange = [&](std::size_t lo, std::size_t hi) {
      // Reuse the traversal stack across every query in this chunk. Tree depth stays below
      // log2(n / LeafSize) + a few for spilled internal pushes, so kDefaultStackReserve rarely
      // grows past its initial capacity; clearing keeps the allocation alive between queries.
      std::vector<KDTreeNode *> stack;
      stack.reserve(kDefaultStackReserve);
      const std::size_t adjReserveFloor =
          std::min(n, (m_dim == kWideAdjReserveDim) ? kWideAdjReserveFloor : kAdjReserveFloor);
      const T *reordered = m_points_reordered.data();
      std::vector<std::pair<std::int32_t, std::int32_t>> &edges =
          workerEdges[math::Pool::workerIndex()];
      const bool useSoa = !m_leafSoa.empty();
      for (std::size_t k = lo; k < hi; ++k) {
        // Walk in tree-build order so consecutive queries share tree paths and keep the visited
        // nodes warm in cache; m_indices[k] maps the reordered row back to its original slot.
        const T *qp = reordered + (k * m_dim);
        const std::size_t rowIdx = m_indices[k];
        std::vector<std::int32_t> &row = out.rows[rowIdx];
        // Each row is filled by exactly this query. Seed a reserve floor so the first survivors do
        // not walk the vector-doubling reallocation cascade from a zero-capacity start.
        row.reserve(adjReserveFloor);
 
        // Clique-leaf shortcut: a swallowed all-core leaf contributes its whole population to
        // the degree and one representative edge to the component build. Taking the shortcut
        // proves this point core (its final degree can only exceed the running guard), so the
        // thinned row stays within the core-row contract; a point that never clears the guard
        // scans normally and keeps a complete row.
        std::size_t bulkDegree = 0;
        stack.clear();
        stack.push_back(m_root);
        while (!stack.empty()) {
          KDTreeNode *node = stack.back();
          stack.pop_back();
          if (node == nullptr) {
            continue;
          }
          const auto [bmin, bmax] = nodeBounds(node);
          if (math::pointAabbGapSq(qp, bmin, bmax) > radius_sq) {
            continue;
          }
          if (node->m_left == nullptr && node->m_right == nullptr) {
            const std::size_t base = node->m_index;
            const std::size_t count = node->m_dim;
            if (leafAllCore[node->m_id] != 0 && row.size() + bulkDegree + count >= minPts &&
                math::pointAabbFarthestSq(qp, bmin, bmax) <= radius_sq) {
              bulkDegree += count;
              const auto rep = static_cast<std::int32_t>(m_indices[base]);
              const auto self = static_cast<std::int32_t>(rowIdx);
              if (rep != self) {
                edges.emplace_back(self, rep);
              }
              continue;
            }
            const T *leafPts = reordered + (base * m_dim);
            auto emit = [&](std::size_t i) noexcept {
              row.push_back(static_cast<std::int32_t>(m_indices[base + i]));
            };
            if (useSoa) {
              math::detail::radiusScanSoa(qp, m_leafSoa.data() + (base * m_dim), count, m_dim,
                                          radius_sq, emit);
            } else {
              math::detail::radiusScan(qp, leafPts, count, m_dim, radius_sq, emit);
            }
            continue;
          }
          const std::size_t pivotSlot = node->m_index;
          const T *pivotRow = reordered + (pivotSlot * m_dim);
          if (math::detail::sqEuclideanRowPtr(qp, pivotRow, m_dim) <= radius_sq) {
            row.push_back(static_cast<std::int32_t>(m_indices[pivotSlot]));
          }
          stack.push_back(node->m_left);
          stack.push_back(node->m_right);
        }
        out.isCore[rowIdx] =
            (row.size() + bulkDegree >= minPts) ? std::uint8_t{1} : std::uint8_t{0};
      }
    };
 
    if (pool.shouldParallelize(n, 4, 2)) {
      // Oversubscribe blocks so dynamic stealing balances the skewed per-query degree; one block
      // per worker lets a dense-neighbourhood region gate the join while the rest idle. Each query
      // is a full tree walk, so the block floor is finer than the row-light default to keep enough
      // blocks per worker at small n.
      pool.parallelForBlocks<citor::HintsDefaults>(
          std::size_t{0}, n, pool.stealBlocks(n, 64),
          [&](std::size_t lo, std::size_t hi) { runRange(lo, hi); });
    } else {
      runRange(0, n);
    }
    for (const auto &edges : workerEdges) {
      out.extraEdges.insert(out.extraEdges.end(), edges.begin(), edges.end());
    }
    return out;
  }

  [[nodiscard]] std::pair<NDArray<std::int32_t, 2>, NDArray<T, 2>> knnQuery(std::int32_t k,
                                                                            math::Pool pool) const {
    const std::size_t n = m_points.dim(0);
    CLUSTERING_ALWAYS_ASSERT(k >= 1);
    CLUSTERING_ALWAYS_ASSERT(std::cmp_less(k, n));
 
    const auto kSz = static_cast<std::size_t>(k);
    NDArray<std::int32_t, 2> indices({n, kSz});
    NDArray<T, 2> sqDists({n, kSz});
 
    auto runRange = [&](std::size_t lo, std::size_t hi) {
      // Reuse the traversal stack and top-k tracker across every query in this chunk; both are
      // reset at the head of each per-point walk.
      std::vector<KDTreeNode *> stack;
      stack.reserve(kDefaultStackReserve);
      math::detail::TopKNeighbors<T, std::int32_t> topK(kSz);
      const T *sourceData = m_points.data();
      std::int32_t *idxOut = indices.data();
      T *distOut = sqDists.data();
      for (std::size_t i = lo; i < hi; ++i) {
        const auto iOriginal = static_cast<std::int32_t>(i);
        const T *qp = sourceData + (i * m_dim);
        topK.clear();
        knnQueryImpl(m_root, qp, iOriginal, topK, stack);
        topK.drainAscending(distOut + (i * kSz), idxOut + (i * kSz));
      }
    };
 
    if (pool.shouldParallelize(n, 4, 2)) {
      pool.parallelForBlocks<citor::HintsDefaults>(
          std::size_t{0}, n, std::size_t{0},
          [&](std::size_t lo, std::size_t hi) { runRange(lo, hi); });
    } else {
      runRange(0, n);
    }
    return {std::move(indices), std::move(sqDists)};
  }

  [[nodiscard]] std::pair<std::span<const T>, std::span<const T>>
  nodeBounds(const KDTreeNode *node) const noexcept {
    assert(node != nullptr && "KDTree::nodeBounds on null node");
    const std::size_t base = static_cast<std::size_t>(node->m_id) * 2 * m_dim;
    const T *bounds = m_nodeBounds.data() + base;
    return {std::span<const T>(bounds, m_dim), std::span<const T>(bounds + m_dim, m_dim)};
  }

  [[nodiscard]] std::span<const std::size_t> indexPermutation() const noexcept {
    return {m_indices.data(), m_indices.size()};
  }

  [[nodiscard]] std::span<const T> reorderedPoints() const noexcept {
    return {m_points_reordered.data(), m_points_reordered.size()};
  }

  [[nodiscard]] std::size_t nodeCount() const noexcept {
    return static_cast<std::size_t>(m_nextNodeId);
  }

  [[nodiscard]] const KDTreeNode *root() const noexcept { return m_root; }

  [[nodiscard]] std::size_t dim() const noexcept { return m_dim; }

  ~KDTree() {
    if (m_allocator.isDeallocSupported()) {
      doRecDealloc(m_root);
    }
  }
 
private:
  static size_t calculatePoolSize(size_t numPoints) {
    if (numPoints == 0) {
      return 0;
    }
    // With leaf-size tuning, the number of nodes is at most numPoints
    // (much less than the original 2*numPoints-1, since leaf nodes batch
    // up to LeafSize points each).
    return numPoints;
  }

  static void doRecDealloc(KDTreeNode *node) {
    if (node == nullptr) {
      return;
    }
 
    doRecDealloc(node->m_left);
    doRecDealloc(node->m_right);
 
    delete node;
  }

  static constexpr std::size_t kParallelBuildFloor = 512;

  static std::size_t nodeCountFor(std::size_t m) noexcept {
    if (m == 0) {
      return 0;
    }
    if (m <= LeafSize) {
      return 1;
    }
    const std::size_t mLeft = (m - 1) / 2;
    return 1 + nodeCountFor(mLeft) + nodeCountFor(m - 1 - mLeft);
  }

  KDTreeNode *buildAt(std::size_t start, std::size_t end, std::size_t depth, KDTreeNode *arena,
                      std::uint32_t idBase, math::Pool pool) {
    if (start >= end) {
      return nullptr;
    }
 
    // Leaf node: store range into m_indices, brute-force at query time
    if (end - start <= LeafSize) {
      *arena = {.m_index = start,     // offset into m_indices
                .m_dim = end - start, // count of points
                .m_left = nullptr,
                .m_right = nullptr,
                .m_id = idBase};
      return arena;
    }
 
    // Internal node: split on median
    const std::size_t dim = depth % m_points.dim(1);
    const std::size_t median = start + (((end - start) - 1) / 2);
 
    using diff_t = std::vector<std::size_t>::difference_type;
    std::nth_element(m_indices.begin() + static_cast<diff_t>(start),
                     m_indices.begin() + static_cast<diff_t>(median),
                     m_indices.begin() + static_cast<diff_t>(end),
                     [this, dim](std::size_t lhs, std::size_t rhs) {
                       return m_points[lhs][dim] < m_points[rhs][dim];
                     });
 
    const std::size_t nLeft = nodeCountFor(median - start);
    const std::size_t nRight = nodeCountFor(end - median - 1);
    KDTreeNode *left = nullptr;
    KDTreeNode *right = nullptr;
    auto buildLeft = [&] { left = buildAt(start, median, depth + 1, arena + 1, idBase, pool); };
    auto buildRight = [&] {
      right = buildAt(median + 1, end, depth + 1, arena + 1 + nLeft,
                      idBase + static_cast<std::uint32_t>(nLeft), pool);
    };
    if (end - start > kParallelBuildFloor && pool.pool != nullptr) {
      pool.forkJoin2(buildLeft, buildRight);
    } else {
      buildLeft();
      buildRight();
    }
 
    *arena = {.m_index = median, // reordered slot of the pivot
              .m_dim = dim,
              .m_left = left,
              .m_right = right,
              .m_id = idBase + static_cast<std::uint32_t>(nLeft + nRight)};
 
    return arena;
  }

  static constexpr std::size_t kBlockQueryDimFloor = 4;

  void blockQuery(T radius_sq, std::size_t minPts, math::Pool pool,
                  index::CoreAdjacency &out) const {
    const std::size_t n = m_points.dim(0);
    const std::size_t totalNodes = nodeCount();
    const KDTreeNode *arenaNodes = m_root;
    std::vector<const KDTreeNode *> leaves;
    std::vector<const KDTreeNode *> pivots;
    leaves.reserve(totalNodes);
    for (std::size_t nodeSlot = 0; nodeSlot < totalNodes; ++nodeSlot) {
      const KDTreeNode &node = arenaNodes[nodeSlot];
      if (node.m_left == nullptr && node.m_right == nullptr) {
        leaves.push_back(&node);
      } else {
        pivots.push_back(&node);
      }
    }
 
    const T *reordered = m_points_reordered.data();
    const std::size_t adjReserveFloor =
        std::min(n, (m_dim == kWideAdjReserveDim) ? kWideAdjReserveFloor : kAdjReserveFloor);
    const bool useSoa = !m_leafSoa.empty();
 
    auto runLeafRange = [&](std::size_t lo, std::size_t hi) {
      std::vector<const KDTreeNode *> stack;
      stack.reserve(kDefaultStackReserve);
      for (std::size_t li = lo; li < hi; ++li) {
        const KDTreeNode &source = *leaves[li];
        const std::size_t base = source.m_index;
        const std::size_t count = source.m_dim;
        const auto [srcMin, srcMax] = nodeBounds(&source);
        for (std::size_t i = 0; i < count; ++i) {
          out.rows[m_indices[base + i]].reserve(adjReserveFloor);
        }
 
        stack.clear();
        stack.push_back(m_root);
        while (!stack.empty()) {
          const KDTreeNode *node = stack.back();
          stack.pop_back();
          if (node == nullptr) {
            continue;
          }
          const auto [nodeMin, nodeMax] = nodeBounds(node);
          if (math::aabbAabbGapSq(srcMin, srcMax, nodeMin, nodeMax) > radius_sq) {
            continue;
          }
          if (node->m_left == nullptr && node->m_right == nullptr) {
            const std::size_t targetBase = node->m_index;
            const std::size_t targetCount = node->m_dim;
            const T *targetPts = reordered + (targetBase * m_dim);
            const T *targetSoa = useSoa ? m_leafSoa.data() + (targetBase * m_dim) : nullptr;
            std::size_t i = 0;
            if (useSoa) {
              // Paired sources share each target-column load and one call setup.
              for (; i + 2 <= count; i += 2) {
                std::vector<std::int32_t> &row0 = out.rows[m_indices[base + i]];
                std::vector<std::int32_t> &row1 = out.rows[m_indices[base + i + 1]];
                math::detail::radiusScanSoaPair(
                    reordered + ((base + i) * m_dim), reordered + ((base + i + 1) * m_dim),
                    targetSoa, targetCount, m_dim, radius_sq,
                    [&](std::size_t j) noexcept {
                      row0.push_back(static_cast<std::int32_t>(m_indices[targetBase + j]));
                    },
                    [&](std::size_t j) noexcept {
                      row1.push_back(static_cast<std::int32_t>(m_indices[targetBase + j]));
                    });
              }
            }
            for (; i < count; ++i) {
              const T *qp = reordered + ((base + i) * m_dim);
              std::vector<std::int32_t> &row = out.rows[m_indices[base + i]];
              auto emit = [&](std::size_t j) noexcept {
                row.push_back(static_cast<std::int32_t>(m_indices[targetBase + j]));
              };
              if (useSoa) {
                math::detail::radiusScanSoa(qp, targetSoa, targetCount, m_dim, radius_sq, emit);
              } else {
                math::detail::radiusScan(qp, targetPts, targetCount, m_dim, radius_sq, emit);
              }
            }
            continue;
          }
          const T *pivotRow = reordered + (node->m_index * m_dim);
          const auto pivotIdx = static_cast<std::int32_t>(m_indices[node->m_index]);
          for (std::size_t i = 0; i < count; ++i) {
            const T *qp = reordered + ((base + i) * m_dim);
            if (math::detail::sqEuclideanRowPtr(qp, pivotRow, m_dim) <= radius_sq) {
              out.rows[m_indices[base + i]].push_back(pivotIdx);
            }
          }
          stack.push_back(node->m_left);
          stack.push_back(node->m_right);
        }
 
        for (std::size_t i = 0; i < count; ++i) {
          const std::size_t rowIdx = m_indices[base + i];
          out.isCore[rowIdx] =
              (out.rows[rowIdx].size() >= minPts) ? std::uint8_t{1} : std::uint8_t{0};
        }
      }
    };
 
    auto runPivotRange = [&](std::size_t lo, std::size_t hi) {
      std::vector<KDTreeNode *> stack;
      stack.reserve(kDefaultStackReserve);
      for (std::size_t pi = lo; pi < hi; ++pi) {
        const std::size_t slot = pivots[pi]->m_index;
        const T *qp = reordered + (slot * m_dim);
        const std::size_t rowIdx = m_indices[slot];
        std::vector<std::int32_t> &row = out.rows[rowIdx];
        row.reserve(adjReserveFloor);
        queryImpl(m_root, qp, radius_sq, row, stack, /*limit=*/-1);
        out.isCore[rowIdx] = (row.size() >= minPts) ? std::uint8_t{1} : std::uint8_t{0};
      }
    };
 
    if (pool.shouldParallelize(n, 4, 2)) {
      pool.parallelForBlocks<citor::HintsDefaults>(
          std::size_t{0}, leaves.size(), pool.stealBlocks(leaves.size(), 1), runLeafRange);
      pool.parallelForBlocks(std::size_t{0}, pivots.size(), std::size_t{0}, runPivotRange);
    } else {
      runLeafRange(0, leaves.size());
      runPivotRange(0, pivots.size());
    }
  }

  template <class OutIdx>
  void queryImpl(KDTreeNode *root, const T *qp, T radius_sq, std::vector<OutIdx> &indices,
                 std::vector<KDTreeNode *> &stack, std::int64_t limit = -1) const {
    if (root == nullptr) {
      return;
    }
 
    stack.clear();
    stack.push_back(root);
 
    const T *reorderedBase = m_points_reordered.data();
 
    while (!stack.empty()) {
      const KDTreeNode *node = stack.back();
      stack.pop_back();
 
      if (node == nullptr) {
        continue;
      }
 
      if (limit != -1 && indices.size() == static_cast<std::size_t>(limit)) {
        break;
      }
 
      // Leaf node: brute-force all points in the range. @c m_points_reordered lays points in
      // tree-build order, so a leaf's entries are a contiguous @c count x d block -- one or
      // two cache lines at small @c d. @ref math::detail::radiusScan dispatches into a SIMD
      // kernel when @c d matches a batched width (today @c f32, `d == 2`) and otherwise
      // falls back to the scalar @ref sqEuclideanRowPtr per-row primitive.
      if (node->m_left == nullptr && node->m_right == nullptr) {
        const std::size_t base = node->m_index;
        const std::size_t count = node->m_dim;
        const T *leafPts = reorderedBase + (base * m_dim);
        // The SoA leaf copy exists only when @c d clears @ref kSoaLeafDimFloor; it is built lazily
        // on the first radius query, and when present scans without a horizontal sum; otherwise the
        // AoS scan runs.
        const bool useSoa = !m_leafSoa.empty();
        const T *leafSoa = useSoa ? m_leafSoa.data() + (base * m_dim) : nullptr;
        const bool isBounded = (limit != -1);
        if (!isBounded) {
          // Unbounded radius query (the DBSCAN adjacency sweep): every survivor is kept, so the
          // leaf-scan emit skips the per-point capacity test the bounded path needs.
          auto emit = [&](std::size_t i) noexcept {
            indices.push_back(static_cast<OutIdx>(m_indices[base + i]));
          };
          if (useSoa) {
            math::detail::radiusScanSoa(qp, leafSoa, count, m_dim, radius_sq, emit);
          } else {
            math::detail::radiusScan(qp, leafPts, count, m_dim, radius_sq, emit);
          }
          continue;
        }
        const auto cap = static_cast<std::size_t>(limit);
        auto emit = [&](std::size_t i) noexcept {
          if (indices.size() >= cap) {
            return;
          }
          indices.push_back(static_cast<OutIdx>(m_indices[base + i]));
        };
        if (useSoa) {
          math::detail::radiusScanSoa(qp, leafSoa, count, m_dim, radius_sq, emit);
        } else {
          math::detail::radiusScan(qp, leafPts, count, m_dim, radius_sq, emit);
        }
        if (indices.size() >= cap) {
          break;
        }
        continue;
      }
 
      // Internal node: check the split point and traverse children. `node->m_index` is the
      // pivot's slot in @c m_points_reordered, so the pivot's row lives at
      // @c reorderedBase + slot * m_dim. The original point index (needed for @c indices)
      // comes from `m_indices[slot]`.
      const std::size_t pivotSlot = node->m_index;
      const std::size_t splitDim = node->m_dim;
      const T *pivotRow = reorderedBase + (pivotSlot * m_dim);
      const T dist_sq = math::detail::sqEuclideanRowPtr(qp, pivotRow, m_dim);
      if (dist_sq <= radius_sq) {
        indices.push_back(static_cast<OutIdx>(m_indices[pivotSlot]));
      }
 
      const T pivotCoord = pivotRow[splitDim];
      const T diff = qp[splitDim] - pivotCoord;
      // Range queries scan the near side unconditionally and the far side only when the split
      // plane lies within the radius. Discovery order does not matter here -- every matching leaf
      // is scanned and the adjacency is order independent -- so the symmetric near / far branches
      // collapse into two straight-line conditional pushes. `diff < 0` selects which child is the
      // near side; `farWithinRadius` admits the other when the plane is reachable.
      const bool farWithinRadius = diff * diff <= radius_sq;
      if (node->m_left != nullptr && (diff < 0 || farWithinRadius)) {
        stack.push_back(node->m_left);
      }
      if (node->m_right != nullptr && (diff >= 0 || farWithinRadius)) {
        stack.push_back(node->m_right);
      }
    }
  }

  void knnQueryImpl(KDTreeNode *root, const T *qp, std::int32_t selfIndex,
                    math::detail::TopKNeighbors<T, std::int32_t> &topK,
                    std::vector<KDTreeNode *> &stack) const {
    if (root == nullptr) {
      return;
    }
 
    stack.clear();
    stack.push_back(root);
 
    const T *reorderedBase = m_points_reordered.data();
 
    while (!stack.empty()) {
      const KDTreeNode *node = stack.back();
      stack.pop_back();
 
      if (node == nullptr) {
        continue;
      }
 
      // Current pruning bound. Until the tracker fills, accept everything; once full, the
      // retained worst-key serves as the upper bound.
      const T bound = topK.boundKey();
 
      // AABB gap prune: the minimum possible squared distance from the query to any point
      // under @c node is @c pointAabbGapSq against the subtree's bounding box. The parent's
      // single-axis prune is a strictly weaker lower bound; at d>=4 the full-AABB gap skips
      // subtrees the axis prune lets through, which at d=8 is the dominant share of the kNN
      // walker's remaining internal-node visits.
      if (bound != std::numeric_limits<T>::max()) {
        auto [nmin, nmax] = this->nodeBounds(node);
        const T gapSq = math::pointAabbGapSq<T>(qp, nmin, nmax);
        if (gapSq >= bound) {
          continue;
        }
      }
 
      if (node->m_left == nullptr && node->m_right == nullptr) {
        // Leaf: walk the contiguous count x d block and admit each candidate. Self-exclusion
        // is by original index so a candidate is skipped iff it is the query row. Distances
        // are computed in blocks of four so the horizontal-sum epilogue is shared across
        // neighbours; the top-k admit then runs over the precomputed dsq vector.
        const std::size_t base = node->m_index;
        const std::size_t count = node->m_dim;
        const T *leafPts = reorderedBase + (base * m_dim);
        std::array<T, LeafSize> dsqBuf{};
        math::detail::sqDistancesAosBlock<T>(qp, leafPts, count, m_dim, dsqBuf.data());
        // Batch-level prune: if every distance in the leaf exceeds the current worst retained
        // bound, no candidate can enter the tracker. Scan for the minimum once; the common
        // case at LeafSize=64 with k in [4, 32] is that every distance falls above the bound
        // and the per-entry admit loop (self-exclude cmp, index load, topK.push) is skipped
        // entirely. The scan runs scalar because LeafSize is known at compile time and the
        // auto-vectoriser produces a tight min-reduce with no hsum epilogue of its own.
        if (topK.full()) {
          T minDsq = dsqBuf[0];
          for (std::size_t i = 1; i < count; ++i) {
            if (dsqBuf[i] < minDsq) {
              minDsq = dsqBuf[i];
            }
          }
          if (minDsq >= bound) {
            continue;
          }
        }
        for (std::size_t i = 0; i < count; ++i) {
          const auto pointIdx = static_cast<std::int32_t>(m_indices[base + i]);
          if (pointIdx == selfIndex) {
            continue;
          }
          topK.push(dsqBuf[i], pointIdx);
        }
        continue;
      }
 
      // Internal: test the pivot, then descend near child first so the bound tightens before
      // the far-child prune test. Compute the split-axis delta squared first so both the pivot
      // admit and the descend-far gate can share it: the full d-wide pivot distance is at least
      // the split-axis contribution, so a larger delta squared proves the pivot cannot enter
      // the retained set and the d-wide hsum is skipped entirely.
      const std::size_t pivotSlot = node->m_index;
      const std::size_t splitDim = node->m_dim;
      const T *pivotRow = reorderedBase + (pivotSlot * m_dim);
      const auto pivotIdx = static_cast<std::int32_t>(m_indices[pivotSlot]);
      const T pivotCoord = pivotRow[splitDim];
      const T diff = qp[splitDim] - pivotCoord;
      const T diffSq = diff * diff;
      if (pivotIdx != selfIndex && diffSq <= bound) {
        const T dist_sq = math::detail::sqEuclideanRowPtr(qp, pivotRow, m_dim);
        topK.push(dist_sq, pivotIdx);
      }
      // DFS descends near-first, far-second. Stack is LIFO, so push far before near.
      if (diff < 0) {
        if (diffSq <= bound && node->m_right != nullptr) {
          stack.push_back(node->m_right);
        }
        if (node->m_left != nullptr) {
          stack.push_back(node->m_left);
        }
      } else {
        if (diffSq <= bound && node->m_left != nullptr) {
          stack.push_back(node->m_left);
        }
        if (node->m_right != nullptr) {
          stack.push_back(node->m_right);
        }
      }
    }
  }

  void ensureLeafSoa(math::Pool pool = {}) const {
    if (m_dim < kSoaLeafDimFloor || m_root == nullptr || !m_leafSoa.empty()) {
      return;
    }
    m_leafSoa.resize(m_points_reordered.size());
    // Leaves transpose disjoint slices, so the pass fans out over the contiguous node arena.
    const KDTreeNode *arenaNodes = m_root;
    pool.parallelForBlocks(std::size_t{0}, nodeCount(), std::size_t{0},
                           [&](std::size_t lo, std::size_t hi) {
                             for (std::size_t nodeSlot = lo; nodeSlot < hi; ++nodeSlot) {
                               const KDTreeNode &node = arenaNodes[nodeSlot];
                               if (node.m_left == nullptr && node.m_right == nullptr) {
                                 transposeLeafSoa(node);
                               }
                             }
                           });
  }

  void transposeLeafSoa(const KDTreeNode &leaf) const noexcept {
    const std::size_t base = leaf.m_index;
    const std::size_t count = leaf.m_dim;
    const T *aos = m_points_reordered.data() + (base * m_dim);
    T *soa = m_leafSoa.data() + (base * m_dim);
    for (std::size_t p = 0; p < count; ++p) {
      for (std::size_t f = 0; f < m_dim; ++f) {
        soa[(f * count) + p] = aos[(p * m_dim) + f];
      }
    }
  }

  void populateLeafBounds(const KDTreeNode &node) noexcept {
    T *minOut = m_nodeBounds.data() + (static_cast<std::size_t>(node.m_id) * 2 * m_dim);
    T *maxOut = minOut + m_dim;
    const std::size_t base = node.m_index;
    const std::size_t count = node.m_dim;
    const T *leafPts = m_points_reordered.data() + (base * m_dim);
    for (std::size_t j = 0; j < m_dim; ++j) {
      minOut[j] = leafPts[j];
      maxOut[j] = leafPts[j];
    }
    for (std::size_t i = 1; i < count; ++i) {
      const T *row = leafPts + (i * m_dim);
      for (std::size_t j = 0; j < m_dim; ++j) {
        if (row[j] < minOut[j]) {
          minOut[j] = row[j];
        }
        if (row[j] > maxOut[j]) {
          maxOut[j] = row[j];
        }
      }
    }
  }

  void unionChildBounds(const KDTreeNode &node) noexcept {
    T *minOut = m_nodeBounds.data() + (static_cast<std::size_t>(node.m_id) * 2 * m_dim);
    T *maxOut = minOut + m_dim;
 
    const KDTreeNode *const seed = (node.m_left != nullptr) ? node.m_left : node.m_right;
    const T *seedMin = m_nodeBounds.data() + (static_cast<std::size_t>(seed->m_id) * 2 * m_dim);
    const T *seedMax = seedMin + m_dim;
    for (std::size_t j = 0; j < m_dim; ++j) {
      minOut[j] = seedMin[j];
      maxOut[j] = seedMax[j];
    }
 
    if (node.m_left != nullptr && node.m_right != nullptr) {
      const T *otherMin =
          m_nodeBounds.data() + (static_cast<std::size_t>(node.m_right->m_id) * 2 * m_dim);
      const T *otherMax = otherMin + m_dim;
      for (std::size_t j = 0; j < m_dim; ++j) {
        if (otherMin[j] < minOut[j]) {
          minOut[j] = otherMin[j];
        }
        if (otherMax[j] > maxOut[j]) {
          maxOut[j] = otherMax[j];
        }
      }
    }
 
    // Union-in the pivot's own row so the internal-node box encloses the pivot point too.
    const std::size_t pivotSlot = node.m_index;
    const T *pivotRow = m_points_reordered.data() + (pivotSlot * m_dim);
    for (std::size_t j = 0; j < m_dim; ++j) {
      if (pivotRow[j] < minOut[j]) {
        minOut[j] = pivotRow[j];
      }
      if (pivotRow[j] > maxOut[j]) {
        maxOut[j] = pivotRow[j];
      }
    }
  }

  static constexpr std::size_t kDefaultStackReserve = 64;

  static constexpr std::size_t kAdjReserveFloor = 8;
  static constexpr std::size_t kWideAdjReserveDim = 8;
  static constexpr std::size_t kWideAdjReserveFloor = 128;

  static constexpr std::size_t kSoaLeafDimFloor = 2;
 
  AllocT m_allocator; 
 
  KDTreeNode *m_root = nullptr;       
  const NDArray<T, 2> &m_points;      
  std::size_t m_dim = 0;              
  std::vector<std::size_t> m_indices; 
  std::vector<T> m_points_reordered;  
  mutable std::vector<T> m_leafSoa;
  std::uint32_t m_nextNodeId = 0;
  std::vector<T> m_nodeBounds;
};
 
} // namespace clustering