Adding robust fitting option to the FitPlane (Marian)
[u/mrichter/AliRoot.git] / STAT / TKDTree.h
1 #ifndef ROOT_TKDTree
2 #define ROOT_TKDTree
3
4 #ifndef ROOT_TObject
5 #include "TObject.h"
6 #endif
7
8 #include "TMath.h"
9
10
11 template <typename Index, typename Value> class TKDTree : public TObject
12 {
13 public:
14         enum{
15                 kDimMax = 6
16         };
17
18         TKDTree();
19         TKDTree(Index npoints, Index ndim, UInt_t bsize, Value **data);
20         ~TKDTree();
21         
22         // getters
23         Index*  GetPointsIndexes(Int_t node) const {
24                 if(node < fNnodes) return 0x0;
25                 Int_t offset = (node >= fCrossNode) ? (node-fCrossNode)*fBucketSize : fOffset+(node-fNnodes)*fBucketSize;
26                 return &fIndPoints[offset];
27         }
28         UChar_t GetNodeAxis(Int_t id) const {return (id < 0 || id >= fNnodes) ? 0 : fAxis[id];}
29         Value   GetNodeValue(Int_t id) const {return (id < 0 || id >= fNnodes) ? 0 : fValue[id];}
30         Int_t   GetNNodes() const {return fNnodes;}
31         Value*  GetBoundaries();
32         Value*  GetBoundary(const Int_t node);
33         static  Int_t   GetIndex(Int_t row, Int_t collumn){return collumn+(1<<row);}
34         static void GetCoord(Int_t index, Int_t &row, Int_t &collumn){for (row=0; index>=(16<<row);row+=4); for (; index>=(2<<row);row++);collumn= index-(1<<row);};
35                 Bool_t  FindNearestNeighbors(const Value *point, const Int_t kNN, Index *&i, Value *&d);
36                 Index   FindNode(const Value * point);
37                 void    FindPoint(Value * point, Index &index, Int_t &iter);
38                 void    FindInRangeA(Value * point, Value * delta, Index *res , Index &npoints,Index & iter, Int_t bnode);
39                 void    FindInRangeB(Value * point, Value * delta, Index *res , Index &npoints,Index & iter, Int_t bnode);
40         void    FindBNodeA(Value * point, Value * delta, Int_t &inode);
41         Bool_t  IsTerminal(Index inode) const {return (inode>=fNnodes);}
42         Value           KOrdStat(Index ntotal, Value *a, Index k, Index *index) const;
43         void            MakeBoundaries(Value *range = 0x0);
44         void            SetData(Index npoints, Index ndim, UInt_t bsize, Value **data);
45         void            Spread(Index ntotal, Value *a, Index *index, Value &min, Value &max) const;
46
47 protected:
48         void            Build();  // build tree
49                                                                         
50 private:
51         TKDTree(const TKDTree &); // not implemented
52         TKDTree<Index, Value>& operator=(const TKDTree<Index, Value>&); // not implemented
53         void            CookBoundaries(const Int_t node, Bool_t left);
54
55
56 protected:
57         Bool_t  fDataOwner;  //! Toggle ownership of the data
58         Int_t   fNnodes;     // size of node array
59         Index   fNDim;       // number of dimensions
60         Index   fNDimm;      // dummy 2*fNDim
61         Index   fNpoints;    // number of multidimensional points
62         Index   fBucketSize; // limit statistic for nodes 
63         UChar_t *fAxis;      //[fNnodes] nodes cutting axis
64         Value   *fValue;     //[fNnodes] nodes cutting value
65         Value           *fRange;     //[fNDimm] range of data for each dimension
66         Value   **fData;     //! data points
67         Value           *fBoundaries;//! nodes boundaries - check class doc
68
69 // kNN related data
70 protected:
71         Int_t   fkNNdim;     //! current kNN arrays allocated dimension
72         Index   *fkNN;       //! k nearest neighbors indexes
73         Value   *fkNNdist;   //! k nearest neighbors distances
74         Value   *fDistBuffer;//! working space for kNN
75         Index   *fIndBuffer; //! working space for kNN
76         
77 private:
78         Index   *fIndPoints; //! array of points indexes
79         Int_t   fRowT0;      //! smallest terminal row
80         Int_t   fCrossNode;  //! cross node
81         Int_t   fOffset;     //! offset in fIndPoints
82
83         ClassDef(TKDTree, 1)  // KD tree
84 };
85
86
87 typedef TKDTree<Int_t, Double_t> TKDTreeID;
88 typedef TKDTree<Int_t, Float_t> TKDTreeIF;
89
90 //_________________________________________________________________
91 template <typename  Index, typename Value>
92 void TKDTree<Index, Value>::FindBNodeA(Value *point, Value *delta, Int_t &inode){
93   //
94   // find the smallest node covering the full range - start
95   //
96   inode =0; 
97   for (;inode<fNnodes;){
98     if (TMath::Abs(point[fAxis[inode]] - fValue[inode])<delta[fAxis[inode]]) break;
99     inode = (point[fAxis[inode]] < fValue[inode]) ? (inode*2)+1: (inode*2)+2;
100   }
101 }
102
103 //_________________________________________________________________
104 template <typename  Index, typename Value>
105 Value* TKDTree<Index, Value>::GetBoundaries()
106 {
107   // Get the boundaries
108         if(!fBoundaries) MakeBoundaries();
109         return fBoundaries;
110 }
111
112 //_________________________________________________________________
113 template <typename  Index, typename Value>
114 Value* TKDTree<Index, Value>::GetBoundary(const Int_t node)
115 {
116   // Get a boundary
117         if(!fBoundaries) MakeBoundaries();
118         return &fBoundaries[node*2*fNDim];
119 }
120
121 #endif
122