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/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)
      : 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);
    m_root = build(0, n, 0);
    // Materialize points in tree-build order. After @ref build 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.
    m_points_reordered.resize(n * m_dim);
    const T *src = points.data();
    T *dst = m_points_reordered.data();
    for (std::size_t k = 0; k < n; ++k) {
      const std::size_t src_row = m_indices[k];
      const T *s = src + (src_row * 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.
    m_nodeBounds.assign(static_cast<std::size_t>(m_nextNodeId) * 2 * m_dim, T{});
    if (m_root != nullptr) {
      populateBounds(m_root);
    }
  }

  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;
    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]] std::vector<std::vector<std::int32_t>> query(T radius, math::Pool pool) const {
    const std::size_t n = m_points.dim(0);
    std::vector<std::vector<std::int32_t>> adj(n);
    if (n == 0) {
      return adj;
    }
 
    const T radius_sq = radius * radius;
 
    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 T *sourceData = m_points.data();
      for (std::size_t i = lo; i < hi; ++i) {
        // @p i iterates the source ordering; @c sourceData + i*m_dim is the caller's row.
        const T *qp = sourceData + (i * m_dim);
        queryImpl(m_root, qp, radius_sq, adj[i], stack, /*limit=*/-1);
      }
    };
 
    if (pool.shouldParallelize(n, 4, 2) && pool.pool != nullptr) {
      pool.pool
          ->submit_blocks(std::size_t{0}, n,
                          [&](std::size_t lo, std::size_t hi) { runRange(lo, hi); })
          .wait();
    } else {
      runRange(0, n);
    }
    return adj;
  }

  [[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.pool != nullptr) {
      pool.pool
          ->submit_blocks(std::size_t{0}, n,
                          [&](std::size_t lo, std::size_t hi) { runRange(lo, hi); })
          .wait();
    } 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;
  }

  KDTreeNode *build(std::size_t start, std::size_t end, std::size_t depth) {
    if (start >= end) {
      return nullptr;
    }
 
    // Leaf node: store range into m_indices, brute-force at query time
    if (end - start <= LeafSize) {
      KDTreeNode *node = m_allocator.allocate();
      *node = {.m_index = start,     // offset into m_indices
               .m_dim = end - start, // count of points
               .m_left = nullptr,
               .m_right = nullptr,
               .m_id = m_nextNodeId++};
      return node;
    }
 
    // 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];
                     });
 
    KDTreeNode *node = m_allocator.allocate();
    *node = {.m_index = median, // reordered slot of the pivot
             .m_dim = dim,
             .m_left = build(start, median, depth + 1),
             .m_right = build(median + 1, end, depth + 1),
             .m_id = m_nextNodeId++};
 
    return node;
  }

  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);
        const bool isBounded = (limit != -1);
        const std::size_t cap =
            isBounded ? static_cast<std::size_t>(limit) : std::numeric_limits<std::size_t>::max();
        math::detail::radiusScan(qp, leafPts, count, m_dim, radius_sq, [&](std::size_t i) noexcept {
          if (indices.size() >= cap) {
            return;
          }
          indices.push_back(static_cast<OutIdx>(m_indices[base + i]));
        });
        if (isBounded && 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;
      if (diff < 0) {
        if (node->m_left != nullptr) {
          stack.push_back(node->m_left);
        }
        if (diff * diff <= radius_sq && node->m_right != nullptr) {
          stack.push_back(node->m_right);
        }
      } else {
        if (node->m_right != nullptr) {
          stack.push_back(node->m_right);
        }
        if (diff * diff <= radius_sq && node->m_left != nullptr) {
          stack.push_back(node->m_left);
        }
      }
    }
  }

  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 populateBounds(KDTreeNode *node) noexcept {
    T *minOut = m_nodeBounds.data() + (static_cast<std::size_t>(node->m_id) * 2 * m_dim);
    T *maxOut = minOut + m_dim;
 
    if (node->m_left == nullptr && node->m_right == nullptr) {
      // Leaf: seed from the first stored row, then fold the remaining rows into min/max.
      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];
          }
        }
      }
      return;
    }
 
    // Internal: recurse first so children's bounds are ready, then take their union. The tree
    // invariant guarantees every internal node has at least one child -- the build routine only
    // returns @c nullptr when @p start >= @p end, and the internal-node branch is only entered
    // when the range exceeds @c LeafSize, which is `>= 1`.
    if (node->m_left != nullptr) {
      populateBounds(node->m_left);
    }
    if (node->m_right != nullptr) {
      populateBounds(node->m_right);
    }
 
    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;
 
  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;  
  std::uint32_t m_nextNodeId = 0;
  std::vector<T> m_nodeBounds;
};
 
} // namespace clustering