Subversion Repositories gelsvn

Rev

Rev 61 | Rev 125 | Go to most recent revision | Only display areas with differences | Ignore whitespace | Details | Blame | Last modification | View Log | RSS feed

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