Subversion Repositories gelsvn

Rev

Rev 203 | Only display areas with differences | Ignore whitespace | Details | Blame | Last modification | View Log | RSS feed

Rev 203 Rev 307
1
#ifndef __KDTREE_H
1
#ifndef __KDTREE_H
2
#define __KDTREE_H
2
#define __KDTREE_H
3
 
3
 
4
#include <cmath>
4
#include <cmath>
5
#include <iostream>
5
#include <iostream>
6
#include <vector>
6
#include <vector>
7
#include <algorithm>
7
#include <algorithm>
8
#include "CGLA/CGLA.h"
8
#include "CGLA/CGLA.h"
9
#include "CGLA/ArithVec.h"
9
#include "CGLA/ArithVec.h"
10
 
10
 
11
#if (_MSC_VER >= 1200)
11
#if (_MSC_VER >= 1200)
12
#pragma warning (push)
12
#pragma warning (push)
13
#pragma warning (disable: 4018)
13
#pragma warning (disable: 4018)
14
#endif
14
#endif
15
 
15
 
16
namespace Geometry
16
namespace Geometry
17
{
17
{
18
	/** \brief A classic K-D tree. 
18
	/** \brief A classic K-D tree. 
19
 
19
 
20
			A K-D tree is a good data structure for storing points in space
20
			A K-D tree is a good data structure for storing points in space
21
			and for nearest neighbour queries. It is basically a generalized 
21
			and for nearest neighbour queries. It is basically a generalized 
22
			binary tree in K dimensions. */
22
			binary tree in K dimensions. */
23
	template<class KeyT, class ValT>
23
	template<class KeyT, class ValT>
24
	class KDTree
24
	class KDTree
25
	{
25
	{
26
		typedef typename KeyT::ScalarType ScalarType;
26
		typedef typename KeyT::ScalarType ScalarType;
27
		typedef KeyT KeyType;
27
		typedef KeyT KeyType;
28
		typedef std::vector<KeyT> KeyVectorType;
28
		typedef std::vector<KeyT> KeyVectorType;
29
		typedef std::vector<ValT> ValVectorType;
29
		typedef std::vector<ValT> ValVectorType;
30
	
30
	
31
		/// KDNode struct represents node in KD tree
31
		/// KDNode struct represents node in KD tree
32
		struct KDNode
32
		struct KDNode
33
		{
33
		{
34
			KeyT key;
34
			KeyT key;
35
			ValT val;
35
			ValT val;
36
			short dsc;
36
			short dsc;
37
 
37
 
38
			KDNode(): dsc(0) {}
38
			KDNode(): dsc(0) {}
39
 
39
 
40
			KDNode(const KeyT& _key, const ValT& _val):
40
			KDNode(const KeyT& _key, const ValT& _val):
41
				key(_key), val(_val), dsc(-1) {}
41
				key(_key), val(_val), dsc(-1) {}
42
 
42
 
43
			ScalarType dist(const KeyType& p) const 
43
			ScalarType dist(const KeyType& p) const 
44
			{
44
			{
45
				KeyType dist_vec = p;
45
				KeyType dist_vec = p;
46
				dist_vec  -= key;
46
				dist_vec  -= key;
47
				return dot(dist_vec, dist_vec);
47
				return dot(dist_vec, dist_vec);
48
			}
48
			}
49
		};
49
		};
50
 
50
 
51
		typedef std::vector<KDNode> NodeVecType;
51
		typedef std::vector<KDNode> NodeVecType;
52
		NodeVecType init_nodes;
52
		NodeVecType init_nodes;
53
		NodeVecType nodes;
53
		NodeVecType nodes;
54
		bool is_built;
54
		bool is_built;
55
 
55
 
56
		/// The greatest depth of KD tree.
56
		/// The greatest depth of KD tree.
57
		int max_depth;
57
		int max_depth;
58
 
58
 
59
		/// Total number of elements in tree
59
		/// Total number of elements in tree
60
		int elements;
60
		int elements;
61
 
61
 
62
		/** Comp is a class used for comparing two keys. Comp is constructed
62
		/** Comp is a class used for comparing two keys. Comp is constructed
63
				with the discriminator - i.e. the coordinate of the key that is used
63
				with the discriminator - i.e. the coordinate of the key that is used
64
				for comparing keys - Comp objects are passed to the sort algorithm.*/
64
				for comparing keys - Comp objects are passed to the sort algorithm.*/
65
		class Comp
65
		class Comp
66
		{
66
		{
67
			const int dsc;
67
			const int dsc;
68
		public:
68
		public:
69
			Comp(int _dsc): dsc(_dsc) {}
69
			Comp(int _dsc): dsc(_dsc) {}
70
			bool operator()(const KeyType& k0, const KeyType& k1) const
70
			bool operator()(const KeyType& k0, const KeyType& k1) const
71
			{
71
			{
72
				int dim=KeyType::get_dim();
72
				int dim=KeyType::get_dim();
73
				for(int i=0;i<dim;i++)
73
				for(int i=0;i<dim;i++)
74
					{
74
					{
75
						int j=(dsc+i)%dim;
75
						int j=(dsc+i)%dim;
76
						if(k0[j]<k1[j])
76
						if(k0[j]<k1[j])
77
							return true;
77
							return true;
78
						if(k0[j]>k1[j])
78
						if(k0[j]>k1[j])
79
							return false;
79
							return false;
80
					}
80
					}
81
				return false;
81
				return false;
82
			}
82
			}
83
 
83
 
84
			bool operator()(const KDNode& k0, const KDNode& k1) const
84
			bool operator()(const KDNode& k0, const KDNode& k1) const
85
			{
85
			{
86
				return (*this)(k0.key,k1.key);
86
				return (*this)(k0.key,k1.key);
87
			}
87
			}
88
		};
88
		};
89
 
89
 
90
		/// The dimension -- K
90
		/// The dimension -- K
91
		const int DIM;
91
		const int DIM;
92
 
92
 
93
		/** Passed a vector of keys, this function will construct an optimal tree.
93
		/** Passed a vector of keys, this function will construct an optimal tree.
94
				It is called recursively - second argument is level in tree. */
94
				It is called recursively - second argument is level in tree. */
95
		void optimize(int, int, int, int);
95
		void optimize(int, int, int, int);
96
 
96
 
97
		/** Finde nearest neighbour. */
97
		/** Finde nearest neighbour. */
98
		int closest_point_priv(int, const KeyType&, ScalarType&) const;
98
		int closest_point_priv(int, const KeyType&, ScalarType&) const;
99
 
99
 
100
							
100
							
101
		void in_sphere_priv(int n, 
101
		void in_sphere_priv(int n, 
102
												const KeyType& p, 
102
												const KeyType& p, 
103
												const ScalarType& dist,
103
												const ScalarType& dist,
104
												std::vector<KeyT>& keys,
104
												std::vector<KeyT>& keys,
105
												std::vector<ValT>& vals) const;
105
												std::vector<ValT>& vals) const;
106
	
106
	
107
		/** Finds the optimal discriminator. There are more ways, but this 
107
		/** Finds the optimal discriminator. There are more ways, but this 
108
				function traverses the vector and finds out what dimension has
108
				function traverses the vector and finds out what dimension has
109
				the greatest difference between min and max element. That dimension
109
				the greatest difference between min and max element. That dimension
110
				is used for discriminator */
110
				is used for discriminator */
111
		int opt_disc(int,int) const;
111
		int opt_disc(int,int) const;
112
	
112
	
113
	public:
113
	public:
114
 
114
 
115
		/** Build tree from vector of keys passed as argument. */
115
		/** Build tree from vector of keys passed as argument. */
116
		KDTree():
116
		KDTree():
117
			is_built(false), max_depth(0), DIM(KeyType::get_dim()), elements(0)
117
			is_built(false), max_depth(0), DIM(KeyType::get_dim()), elements(0)
118
		{
118
		{
119
		}
119
		}
120
 
120
 
121
		/** Insert a key value pair into the tree. Note that the tree needs to 
121
		/** Insert a key value pair into the tree. Note that the tree needs to 
122
				be built - by calling the build function - before you can search. */
122
				be built - by calling the build function - before you can search. */
123
		void insert(const KeyT& key, const ValT& val)
123
		void insert(const KeyT& key, const ValT& val)
124
		{
124
		{
125
			assert(!is_built);
125
			assert(!is_built);
126
			init_nodes.push_back(KDNode(key,val));
126
			init_nodes.push_back(KDNode(key,val));
127
		}
127
		}
128
 
128
 
129
		/** Build the tree. After this function have been called, it is no longer 
129
		/** Build the tree. After this function have been called, it is no longer 
130
				legal to insert elements, but you can perform searches. */
130
				legal to insert elements, but you can perform searches. */
131
		void build()
131
		void build()
132
		{
132
		{
133
			assert(!is_built);
133
			assert(!is_built);
134
			nodes.resize(init_nodes.size()+1);
134
			nodes.resize(init_nodes.size()+1);
135
			if(init_nodes.size() > 0)	
135
			if(init_nodes.size() > 0)	
136
				optimize(1,0,init_nodes.size(),0);
136
				optimize(1,0,init_nodes.size(),0);
137
			NodeVecType v(0);
137
			NodeVecType v(0);
138
			init_nodes.swap(v);
138
			init_nodes.swap(v);
139
			is_built = true;
139
			is_built = true;
140
		}
140
		}
141
 
141
 
142
		/** Find the key value pair closest to the key given as first 
142
		/** Find the key value pair closest to the key given as first 
143
				argument. The second argument is the maximum search distance.
143
				argument. The second argument is the maximum search distance.
144
				The final two arguments contain the closest key and its 
144
				The final two arguments contain the closest key and its 
145
				associated value upon return. */
145
				associated value upon return. */
146
		bool closest_point(const KeyT& p, float& dist, KeyT&k, ValT&v) const
146
		bool closest_point(const KeyT& p, float& dist, KeyT&k, ValT&v) const
147
		{
147
		{
148
			assert(is_built);
148
			assert(is_built);
149
			float max_sq_dist = CGLA::sqr(dist);
149
			float max_sq_dist = CGLA::sqr(dist);
150
			if(int n = closest_point_priv(1, p, max_sq_dist))
150
			if(int n = closest_point_priv(1, p, max_sq_dist))
151
				{
151
				{
152
					k = nodes[n].key;
152
					k = nodes[n].key;
153
					v = nodes[n].val;
153
					v = nodes[n].val;
154
					dist = std::sqrt(max_sq_dist);
154
					dist = std::sqrt(max_sq_dist);
155
					return true;
155
					return true;
156
				}
156
				}
157
			return false;
157
			return false;
158
		}
158
		}
159
 
159
 
160
		/** Find all the elements within a given radius (second argument) of
160
		/** Find all the elements within a given radius (second argument) of
161
				the key (first argument). The key value pairs inside the sphere are
161
				the key (first argument). The key value pairs inside the sphere are
162
				returned in a pair of vectors passed as the two last arguments. */
162
				returned in a pair of vectors passed as the two last arguments. */
163
		int in_sphere(const KeyType& p, 
163
		int in_sphere(const KeyType& p, 
164
									float dist,
164
									float dist,
165
									std::vector<KeyT>& keys,
165
									std::vector<KeyT>& keys,
166
									std::vector<ValT>& vals) const
166
									std::vector<ValT>& vals) const
167
		{
167
		{
168
			assert(is_built);
168
			assert(is_built);
169
			float max_sq_dist = CGLA::sqr(dist);
169
			float max_sq_dist = CGLA::sqr(dist);
170
			in_sphere_priv(1,p,max_sq_dist,keys,vals);
170
			in_sphere_priv(1,p,max_sq_dist,keys,vals);
171
			return keys.size();
171
			return keys.size();
172
		}
172
		}
173
		
173
		
174
 
174
 
175
	};
175
	};
176
 
176
 
177
	template<class KeyT, class ValT>
177
	template<class KeyT, class ValT>
178
	int KDTree<KeyT,ValT>::opt_disc(int kvec_beg,  
178
	int KDTree<KeyT,ValT>::opt_disc(int kvec_beg,  
179
																	int kvec_end) const 
179
																	int kvec_end) const 
180
	{
180
	{
181
		KeyType vmin = init_nodes[kvec_beg].key;
181
		KeyType vmin = init_nodes[kvec_beg].key;
182
		KeyType vmax = init_nodes[kvec_beg].key;
182
		KeyType vmax = init_nodes[kvec_beg].key;
183
		for(int i=kvec_beg;i<kvec_end;i++)
183
		for(int i=kvec_beg;i<kvec_end;i++)
184
			{
184
			{
185
				vmin = CGLA::v_min(vmin,init_nodes[i].key);
185
				vmin = CGLA::v_min(vmin,init_nodes[i].key);
186
				vmax = CGLA::v_max(vmax,init_nodes[i].key);
186
				vmax = CGLA::v_max(vmax,init_nodes[i].key);
187
			}
187
			}
188
		int od=0;
188
		int od=0;
189
		KeyType ave_v = vmax-vmin;
189
		KeyType ave_v = vmax-vmin;
190
		for(int i=1;i<KeyType::get_dim();i++)
190
		for(int i=1;i<KeyType::get_dim();i++)
191
			if(ave_v[i]>ave_v[od]) od = i;
191
			if(ave_v[i]>ave_v[od]) od = i;
192
		return od;
192
		return od;
193
	} 
193
	} 
194
 
194
 
195
	template<class KeyT, class ValT>
195
	template<class KeyT, class ValT>
196
	void KDTree<KeyT,ValT>::optimize(int cur,
196
	void KDTree<KeyT,ValT>::optimize(int cur,
197
																	 int kvec_beg,  
197
																	 int kvec_beg,  
198
																	 int kvec_end,  
198
																	 int kvec_end,  
199
																	 int level)
199
																	 int level)
200
	{
200
	{
201
		// Assert that we are not inserting beyond capacity.
201
		// Assert that we are not inserting beyond capacity.
202
		assert(cur < nodes.size());
202
		assert(cur < nodes.size());
203
 
203
 
204
		// If there is just a single element, we simply insert.
204
		// If there is just a single element, we simply insert.
205
		if(kvec_beg+1==kvec_end) 
205
		if(kvec_beg+1==kvec_end) 
206
			{
206
			{
207
				max_depth  = std::max(level,max_depth);
207
				max_depth  = std::max(level,max_depth);
208
				nodes[cur] = init_nodes[kvec_beg];
208
				nodes[cur] = init_nodes[kvec_beg];
209
				nodes[cur].dsc = -1;
209
				nodes[cur].dsc = -1;
210
				return;
210
				return;
211
			}
211
			}
212
	
212
	
213
		// Find the axis that best separates the data.
213
		// Find the axis that best separates the data.
214
		int disc = opt_disc(kvec_beg, kvec_end);
214
		int disc = opt_disc(kvec_beg, kvec_end);
215
 
215
 
216
		// Compute the median element. See my document on how to do this
216
		// Compute the median element. See my document on how to do this
217
		// www.imm.dtu.dk/~jab/publications.html
217
		// www.imm.dtu.dk/~jab/publications.html
218
		int N = kvec_end-kvec_beg;
218
		int N = kvec_end-kvec_beg;
219
		int M = 1<< (CGLA::two_to_what_power(N));
219
		int M = 1<< (CGLA::two_to_what_power(N));
220
		int R = N-(M-1);
220
		int R = N-(M-1);
221
		int left_size  = (M-2)/2;
221
		int left_size  = (M-2)/2;
222
		int right_size = (M-2)/2;
222
		int right_size = (M-2)/2;
223
		if(R < M/2)
223
		if(R < M/2)
224
			{
224
			{
225
				left_size += R;
225
				left_size += R;
226
			}
226
			}
227
		else
227
		else
228
			{
228
			{
229
				left_size += M/2;
229
				left_size += M/2;
230
				right_size += R-M/2;
230
				right_size += R-M/2;
231
			}
231
			}
232
 
232
 
233
		int median = kvec_beg + left_size;
233
		int median = kvec_beg + left_size;
234
 
234
 
235
		// Sort elements but use nth_element (which is cheaper) than
235
		// Sort elements but use nth_element (which is cheaper) than
236
		// a sorting algorithm. All elements to the left of the median
236
		// a sorting algorithm. All elements to the left of the median
237
		// will be smaller than or equal the median. All elements to the right
237
		// will be smaller than or equal the median. All elements to the right
238
		// will be greater than or equal to the median.
238
		// will be greater than or equal to the median.
239
		const Comp comp(disc);
239
		const Comp comp(disc);
240
		std::nth_element(&init_nodes[kvec_beg], 
240
		std::nth_element(&init_nodes[kvec_beg], 
241
										 &init_nodes[median], 
241
										 &init_nodes[median], 
242
										 &init_nodes[kvec_end], comp);
242
										 &init_nodes[kvec_end], comp);
243
 
243
 
244
		// Insert the node in the final data structure.
244
		// Insert the node in the final data structure.
245
		nodes[cur] = init_nodes[median];
245
		nodes[cur] = init_nodes[median];
246
		nodes[cur].dsc = disc;
246
		nodes[cur].dsc = disc;
247
 
247
 
248
		// Recursively build left and right tree.
248
		// Recursively build left and right tree.
249
		if(left_size>0)	
249
		if(left_size>0)	
250
			optimize(2*cur, kvec_beg, median,level+1);
250
			optimize(2*cur, kvec_beg, median,level+1);
251
		
251
		
252
		if(right_size>0) 
252
		if(right_size>0) 
253
			optimize(2*cur+1, median+1, kvec_end,level+1);
253
			optimize(2*cur+1, median+1, kvec_end,level+1);
254
	}
254
	}
255
 
255
 
256
	template<class KeyT, class ValT>
256
	template<class KeyT, class ValT>
257
	int KDTree<KeyT,ValT>::closest_point_priv(int n, const KeyType& p, 
257
	int KDTree<KeyT,ValT>::closest_point_priv(int n, const KeyType& p, 
258
																						ScalarType& dist) const
258
																						ScalarType& dist) const
259
	{
259
	{
260
		int ret_node = 0;
260
		int ret_node = 0;
261
		ScalarType this_dist = nodes[n].dist(p);
261
		ScalarType this_dist = nodes[n].dist(p);
262
 
262
 
263
		if(this_dist<dist)
263
		if(this_dist<dist)
264
			{
264
			{
265
				dist = this_dist;
265
				dist = this_dist;
266
				ret_node = n;
266
				ret_node = n;
267
			}
267
			}
268
		if(nodes[n].dsc != -1)
268
		if(nodes[n].dsc != -1)
269
			{
269
			{
270
				int dsc         = nodes[n].dsc;
270
				int dsc         = nodes[n].dsc;
271
				float dsc_dist  = CGLA::sqr(nodes[n].key[dsc]-p[dsc]);
271
				float dsc_dist  = CGLA::sqr(nodes[n].key[dsc]-p[dsc]);
272
				bool left_son   = Comp(dsc)(p,nodes[n].key);
272
				bool left_son   = Comp(dsc)(p,nodes[n].key);
273
 
273
 
274
				if(left_son||dsc_dist<dist)
274
				if(left_son||dsc_dist<dist)
275
					{
275
					{
276
						int left_child = 2*n;
276
						int left_child = 2*n;
277
						if(left_child < nodes.size())
277
						if(left_child < nodes.size())
278
							if(int nl=closest_point_priv(left_child, p, dist))
278
							if(int nl=closest_point_priv(left_child, p, dist))
279
								ret_node = nl;
279
								ret_node = nl;
280
					}
280
					}
281
				if(!left_son||dsc_dist<dist)
281
				if(!left_son||dsc_dist<dist)
282
					{
282
					{
283
						int right_child = 2*n+1;
283
						int right_child = 2*n+1;
284
						if(right_child < nodes.size())
284
						if(right_child < nodes.size())
285
							if(int nr=closest_point_priv(right_child, p, dist))
285
							if(int nr=closest_point_priv(right_child, p, dist))
286
								ret_node = nr;
286
								ret_node = nr;
287
					}
287
					}
288
			}
288
			}
289
		return ret_node;
289
		return ret_node;
290
	}
290
	}
291
 
291
 
292
	template<class KeyT, class ValT>
292
	template<class KeyT, class ValT>
293
	void KDTree<KeyT,ValT>::in_sphere_priv(int n, 
293
	void KDTree<KeyT,ValT>::in_sphere_priv(int n, 
294
																				 const KeyType& p, 
294
																				 const KeyType& p, 
295
																				 const ScalarType& dist,
295
																				 const ScalarType& dist,
296
																				 std::vector<KeyT>& keys,
296
																				 std::vector<KeyT>& keys,
297
																				 std::vector<ValT>& vals) const
297
																				 std::vector<ValT>& vals) const
298
	{
298
	{
299
		ScalarType this_dist = nodes[n].dist(p);
299
		ScalarType this_dist = nodes[n].dist(p);
300
		assert(n<nodes.size());
300
		assert(n<nodes.size());
301
		if(this_dist<dist)
301
		if(this_dist<dist)
302
			{
302
			{
303
				keys.push_back(nodes[n].key);
303
				keys.push_back(nodes[n].key);
304
				vals.push_back(nodes[n].val);
304
				vals.push_back(nodes[n].val);
305
			}
305
			}
306
		if(nodes[n].dsc != -1)
306
		if(nodes[n].dsc != -1)
307
			{
307
			{
308
				const int dsc         = nodes[n].dsc;
308
				const int dsc         = nodes[n].dsc;
309
				const float dsc_dist  = CGLA::sqr(nodes[n].key[dsc]-p[dsc]);
309
				const float dsc_dist  = CGLA::sqr(nodes[n].key[dsc]-p[dsc]);
310
 
310
 
311
				bool left_son = Comp(dsc)(p,nodes[n].key);
311
				bool left_son = Comp(dsc)(p,nodes[n].key);
312
 
312
 
313
				if(left_son||dsc_dist<dist)
313
				if(left_son||dsc_dist<dist)
314
					{
314
					{
315
						int left_child = 2*n;
315
						int left_child = 2*n;
316
						if(left_child < nodes.size())
316
						if(left_child < nodes.size())
317
							in_sphere_priv(left_child, p, dist, keys, vals);
317
							in_sphere_priv(left_child, p, dist, keys, vals);
318
					}
318
					}
319
				if(!left_son||dsc_dist<dist)
319
				if(!left_son||dsc_dist<dist)
320
					{
320
					{
321
						int right_child = 2*n+1;
321
						int right_child = 2*n+1;
322
						if(right_child < nodes.size())
322
						if(right_child < nodes.size())
323
							in_sphere_priv(right_child, p, dist, keys, vals);
323
							in_sphere_priv(right_child, p, dist, keys, vals);
324
					}
324
					}
325
			}
325
			}
326
	}
326
	}
327
}
327
}
328
namespace GEO = Geometry;
328
namespace GEO = Geometry;
329
 
329
 
330
#if (_MSC_VER >= 1200)
330
#if (_MSC_VER >= 1200)
331
#pragma warning (pop)
331
#pragma warning (pop)
332
#endif
332
#endif
333
 
333
 
334
 
334
 
335
#endif
335
#endif
336
 
336