KDTree.hh

Go to the documentation of this file.

#ifndef G_KDTREE2_H
#define G_KDTREE2_H
 
// (c) Matthew B. Kennel, Institute for Nonlinear Science, UCSD (2004)
//
// Licensed under the Academic Free License version 1.1 found in file LICENSE
// with additional provisions in that same file.
 
// Implement a kd tree for fast searching of points in a fixed data base
// in k-dimensional Euclidean space.
 
#include <array>
#include <queue>
#include <vector>
#include <cstddef>
 
namespace Garfield {
 
typedef std::vector<std::vector<double> > KDTreeArray;
 
class KDTreeNode;

 
struct KDTreeResult {
  double dis;  //< square Euclidean distance
  std::size_t idx;  //< index
};

 
class KDTree {
 public:
  // Reference to the underlying data to be included in the tree.
  const KDTreeArray& m_data;
 
  std::size_t m_dim = 3;
  bool sort_results = false;
 
 public:
  KDTree() = delete;
  KDTree(KDTreeArray& data_in);
  ~KDTree();

  void n_nearest(const std::vector<double>& qv, const unsigned int nn,
                 std::vector<KDTreeResult>& result) const;

  void n_nearest_around_point(const unsigned int idx,
                              const unsigned int ndecorrel,
                              const unsigned int nn,
                              std::vector<KDTreeResult>& result) const;

  void r_nearest(const std::vector<double>& qv, const double r2,
                 std::vector<KDTreeResult>& result) const;

  void r_nearest_around_point(const unsigned int idx,
                              const unsigned int ndecorrel, const double r2,
                              std::vector<KDTreeResult>& result) const;
 
  friend class KDTreeNode;
 
 private:
  KDTreeNode* m_root = nullptr;
 
  // Index for the tree leaves. Data in a leaf with bounds [l,u] are
  // in 'data[ind[l],*] to data[ind[u],*]
  std::vector<std::size_t> m_ind;
 
  static constexpr int bucketsize = 12;  // global constant.
 
 private:
  KDTreeNode* build_tree_for_range(int l, int u, KDTreeNode* parent);
  int select_on_coordinate_value(int c, double alpha, int l, int u);
  std::array<double, 2> spread_in_coordinate(const int c, const int l,
                                             const int u) const;
};

 
class KDTreeNode {
 public:
  KDTreeNode(int dim);
  ~KDTreeNode();
 
 private:
  friend class KDTree;
 
  // Dimension to cut.
  std::size_t cut_dim = 0;
  // Cut value.
  double cut_val = 0.;
  double cut_val_left = 0.;
  double cut_val_right = 0.;
  // Extents in index array for searching
  int m_l = 0;
  int m_u = 0;
  // [min,max] of the box enclosing all points
  std::vector<std::array<double, 2> > box;
 
  // Pointers to left and right nodes.
  KDTreeNode* left = nullptr;
  KDTreeNode* right = nullptr;
 
  // Recursive innermost core routines for searching.
  void search_n(const int idx0, const int nd, const unsigned int nn, double& r2,
                const std::vector<double>& qv, const KDTree& tree,
                std::priority_queue<KDTreeResult>& res) const;
  void search_r(const int idx0, const int nd, const double r2,
                const std::vector<double>& qv, const KDTree& tree,
                std::vector<KDTreeResult>& res) const;
 
  // Return true if the bounding box for this node is within the
  // search range around a point.
  bool box_in_search_range(const double r2,
                           const std::vector<double>& qv) const;
 
  // For processing final buckets.
  void process_terminal_node_n(const int idx0, const int nd,
                               const unsigned int nn, double& r2,
                               const std::vector<double>& qv,
                               const KDTree& tree,
                               std::priority_queue<KDTreeResult>& res) const;
  void process_terminal_node_r(const int idx0, const int nd, const double r2,
                               const std::vector<double>& qv,
                               const KDTree& tree,
                               std::vector<KDTreeResult>& res) const;
};
 
}  // namespace Garfield
 
#endif