include/root/TKDTree.h

0001 #ifndef ROOT_TKDTree
0002 #define ROOT_TKDTree
0003
0004 #include "TObject.h"
0005
0006 #include "TMath.h"
0007 #include <vector>
0008
0009 template <typename Index, typename Value> class TKDTree : public TObject
0010 {
0011 public:
0012
0013    TKDTree();
0014    TKDTree(Index npoints, Index ndim, UInt_t bsize);
0015    TKDTree(Index npoints, Index ndim, UInt_t bsize, Value **data);
0016    ~TKDTree() override;
0017
0018    void            Build();  // build the tree
0019
0020    Double_t        Distance(const Value *point, Index ind, Int_t type=2) const;
0021    void            DistanceToNode(const Value *point, Index inode, Value &min, Value &max, Int_t type=2);
0022
0023    // Get indexes of left and right daughter nodes
0024    Int_t   GetLeft(Int_t inode)  const    {return inode*2+1;}
0025    Int_t   GetRight(Int_t inode) const    {return (inode+1)*2;}
0026    Int_t   GetParent(Int_t inode) const  {return (inode-1)/2;}
0027    //
0028    // Other getters
0029    Index*  GetPointsIndexes(Int_t node) const;
0030    void    GetNodePointsIndexes(Int_t node, Int_t &first1, Int_t &last1, Int_t &first2, Int_t &last2) const;
0031    UChar_t GetNodeAxis(Int_t id) const {return (id < 0 || id >= fNNodes) ? 0 : fAxis[id];}
0032    Value   GetNodeValue(Int_t id) const {return (id < 0 || id >= fNNodes) ? 0 : fValue[id];}
0033    Int_t   GetNNodes() const {return fNNodes;}
0034    Int_t   GetTotalNodes() const {return fTotalNodes;}
0035    Value*  GetBoundaries();
0036    Value*  GetBoundariesExact();
0037    Value*  GetBoundary(const Int_t node);
0038    Value*  GetBoundaryExact(const Int_t node);
0039    Index   GetNPoints() { return fNPoints; }
0040    Index   GetNDim()    { return fNDim; }
0041    Index   GetNPointsNode(Int_t node) const;
0042
0043    //Getters for internal variables.
0044    Int_t   GetRowT0() {return fRowT0;}      //! smallest terminal row
0045    Int_t   GetCrossNode() {return fCrossNode;}  //! cross node
0046    Int_t   GetOffset() {return fOffset;}     //! offset in fIndPoints
0047    Index*  GetIndPoints() {return fIndPoints;}
0048    Index   GetBucketSize() {return fBucketSize;}
0049
0050    void    FindNearestNeighbors(const Value *point, Int_t k, Index *ind, Value *dist);
0051    Index   FindNode(const Value * point) const;
0052    void    FindPoint(Value * point, Index &index, Int_t &iter);
0053    void    FindInRange(Value *point, Value range, std::vector<Index> &res);
0054    void    FindBNodeA(Value * point, Value * delta, Int_t &inode);
0055
0056    Bool_t  IsTerminal(Index inode) const {return (inode>=fNNodes);}
0057    Int_t   IsOwner() { return fDataOwner; }
0058    Value   KOrdStat(Index ntotal, Value *a, Index k, Index *index) const;
0059
0060
0061    void    MakeBoundaries(Value *range = nullptr);
0062    void    MakeBoundariesExact();
0063    void    SetData(Index npoints, Index ndim, UInt_t bsize, Value **data);
0064    Int_t   SetData(Index idim, Value *data);
0065    void    SetOwner(Int_t owner) { fDataOwner = owner; }
0066    void    Spread(Index ntotal, Value *a, Index *index, Value &min, Value &max) const;
0067
0068  private:
0069    TKDTree(const TKDTree &); // not implemented
0070    TKDTree<Index, Value>& operator=(const TKDTree<Index, Value>&); // not implemented
0071    void CookBoundaries(const Int_t node, Bool_t left);
0072
0073    void UpdateNearestNeighbors(Index inode, const Value *point, Int_t kNN, Index *ind, Value *dist);
0074    void UpdateRange(Index inode, Value *point, Value range, std::vector<Index> &res);
0075
0076  protected:
0077    Int_t   fDataOwner;  ///<! 0 - not owner, 2 - owner of the pointer array, 1 - owner of the whole 2-d array
0078    Int_t   fNNodes;     ///< size of node array
0079    Int_t   fTotalNodes; ///< total number of nodes (fNNodes + terminal nodes)
0080    Index   fNDim;       ///< number of dimensions
0081    Index   fNDimm;      ///< dummy 2*fNDim
0082    Index   fNPoints;    ///< number of multidimensional points
0083    Index   fBucketSize; ///< size of the terminal nodes
0084    UChar_t *fAxis;      ///<[fNNodes] nodes cutting axis
0085    Value   *fValue;     ///<[fNNodes] nodes cutting value
0086    //
0087    Value   *fRange;     ///<[fNDimm] range of data for each dimension
0088    Value   **fData;     ///<! data points
0089    Value   *fBoundaries;///<! nodes boundaries
0090
0091
0092    Index   *fIndPoints; ///<! array of points indexes
0093    Int_t   fRowT0;      ///<! smallest terminal row - first row that contains terminal nodes
0094    Int_t   fCrossNode;  ///<! cross node - node that begins the last row (with terminal nodes only)
0095    Int_t   fOffset;     ///<! offset in fIndPoints - if there are 2 rows, that contain terminal nodes
0096                         ///<  fOffset returns the index in the fIndPoints array of the first point
0097                         ///<  that belongs to the first node on the second row.
0098
0099
0100    ClassDefOverride(TKDTree, 1)  // KD tree
0101 };
0102
0103
0104 typedef TKDTree<Int_t, Double_t> TKDTreeID;
0105 typedef TKDTree<Int_t, Float_t> TKDTreeIF;
0106
0107 // Test functions:
0108 TKDTreeIF *  TKDTreeTestBuild();
0109
0110 #endif
0111