backported kd tree improvements from monav project: faster with base case 8
This commit is contained in:
parent
b87d6f3c66
commit
171815c9b7
@ -34,18 +34,20 @@ KD Tree coded by Christian Vetter, Monav Project
|
||||
|
||||
namespace KDTree {
|
||||
|
||||
#define KDTREE_BASESIZE (8)
|
||||
|
||||
template< unsigned k, typename T >
|
||||
class BoundingBox {
|
||||
public:
|
||||
BoundingBox() {
|
||||
for ( unsigned dim = 0; dim < k; ++dim ) {
|
||||
min[dim] = std::numeric_limits< T >::min();
|
||||
max[dim] = std::numeric_limits< T >::max();
|
||||
}
|
||||
}
|
||||
BoundingBox() {
|
||||
for ( unsigned dim = 0; dim < k; ++dim ) {
|
||||
min[dim] = std::numeric_limits< T >::min();
|
||||
max[dim] = std::numeric_limits< T >::max();
|
||||
}
|
||||
}
|
||||
|
||||
T min[k];
|
||||
T max[k];
|
||||
T min[k];
|
||||
T max[k];
|
||||
};
|
||||
|
||||
struct NoData {};
|
||||
@ -53,158 +55,174 @@ struct NoData {};
|
||||
template< unsigned k, typename T >
|
||||
class EuclidianMetric {
|
||||
public:
|
||||
double operator() ( const T left[k], const T right[k] ) {
|
||||
double result = 0;
|
||||
for ( unsigned i = 0; i < k; ++i ) {
|
||||
double temp = (double)left[i] - (double)right[i];
|
||||
result += temp * temp;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
double operator() ( const T left[k], const T right[k] ) {
|
||||
double result = 0;
|
||||
for ( unsigned i = 0; i < k; ++i ) {
|
||||
double temp = (double)left[i] - (double)right[i];
|
||||
result += temp * temp;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
double operator() ( const BoundingBox< k, T > &box, const T point[k] ) {
|
||||
T nearest[k];
|
||||
for ( unsigned dim = 0; dim < k; ++dim ) {
|
||||
if ( point[dim] < box.min[dim] )
|
||||
nearest[dim] = box.min[dim];
|
||||
else if ( point[dim] > box.max[dim] )
|
||||
nearest[dim] = box.max[dim];
|
||||
else
|
||||
nearest[dim] = point[dim];
|
||||
}
|
||||
return operator() ( point, nearest );
|
||||
}
|
||||
double operator() ( const BoundingBox< k, T > &box, const T point[k] ) {
|
||||
T nearest[k];
|
||||
for ( unsigned dim = 0; dim < k; ++dim ) {
|
||||
if ( point[dim] < box.min[dim] )
|
||||
nearest[dim] = box.min[dim];
|
||||
else if ( point[dim] > box.max[dim] )
|
||||
nearest[dim] = box.max[dim];
|
||||
else
|
||||
nearest[dim] = point[dim];
|
||||
}
|
||||
return operator() ( point, nearest );
|
||||
}
|
||||
};
|
||||
|
||||
template < unsigned k, typename T, typename Data = NoData, typename Metric = EuclidianMetric< k, T > >
|
||||
class StaticKDTree {
|
||||
public:
|
||||
|
||||
struct InputPoint {
|
||||
T coordinates[k];
|
||||
Data data;
|
||||
};
|
||||
struct InputPoint {
|
||||
T coordinates[k];
|
||||
Data data;
|
||||
bool operator==( const InputPoint& right )
|
||||
{
|
||||
for ( int i = 0; i < k; i++ ) {
|
||||
if ( coordinates[i] != right.coordinates[i] )
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
StaticKDTree( std::vector< InputPoint > * points ){
|
||||
assert( k > 0 );
|
||||
assert ( points->size() > 0 );
|
||||
size = points->size();
|
||||
kdtree = new InputPoint[size];
|
||||
for ( Iterator i = 0; i != size; ++i ) {
|
||||
kdtree[i] = points->at(i);
|
||||
for ( unsigned dim = 0; dim < k; ++dim ) {
|
||||
if ( kdtree[i].coordinates[dim] < boundingBox.min[dim] )
|
||||
boundingBox.min[dim] = kdtree[i].coordinates[dim];
|
||||
if ( kdtree[i].coordinates[dim] > boundingBox.max[dim] )
|
||||
boundingBox.max[dim] = kdtree[i].coordinates[dim];
|
||||
}
|
||||
}
|
||||
std::stack< Tree > s;
|
||||
s.push ( Tree ( 0, size, 0 ) );
|
||||
while ( !s.empty() ) {
|
||||
Tree tree = s.top();
|
||||
s.pop();
|
||||
StaticKDTree( std::vector< InputPoint > * points ){
|
||||
assert( k > 0 );
|
||||
assert ( points->size() > 0 );
|
||||
size = points->size();
|
||||
kdtree = new InputPoint[size];
|
||||
for ( Iterator i = 0; i != size; ++i ) {
|
||||
kdtree[i] = points->at(i);
|
||||
for ( unsigned dim = 0; dim < k; ++dim ) {
|
||||
if ( kdtree[i].coordinates[dim] < boundingBox.min[dim] )
|
||||
boundingBox.min[dim] = kdtree[i].coordinates[dim];
|
||||
if ( kdtree[i].coordinates[dim] > boundingBox.max[dim] )
|
||||
boundingBox.max[dim] = kdtree[i].coordinates[dim];
|
||||
}
|
||||
}
|
||||
std::stack< Tree > s;
|
||||
s.push ( Tree ( 0, size, 0 ) );
|
||||
while ( !s.empty() ) {
|
||||
Tree tree = s.top();
|
||||
s.pop();
|
||||
|
||||
if ( tree.left == tree.right )
|
||||
continue;
|
||||
if ( tree.right - tree.left < KDTREE_BASESIZE )
|
||||
continue;
|
||||
|
||||
Iterator middle = tree.left + ( tree.right - tree.left ) / 2;
|
||||
Iterator middle = tree.left + ( tree.right - tree.left ) / 2;
|
||||
#ifdef _GLIBCXX_PARALLEL
|
||||
__gnu_parallel::nth_element( kdtree + tree.left, kdtree + middle, kdtree + tree.right, Less( tree.dimension ) );
|
||||
#else
|
||||
std::nth_element( kdtree + tree.left, kdtree + middle, kdtree + tree.right, Less( tree.dimension ) );
|
||||
__gnu_parallel::nth_element( kdtree + tree.left, kdtree + middle, kdtree + tree.right, Less( tree.dimension ) );
|
||||
#else
|
||||
std::nth_element( kdtree + tree.left, kdtree + middle, kdtree + tree.right, Less( tree.dimension ) );
|
||||
#endif
|
||||
s.push( Tree( tree.left, middle, ( tree.dimension + 1 ) % k ) );
|
||||
s.push( Tree( middle + 1, tree.right, ( tree.dimension + 1 ) % k ) );
|
||||
}
|
||||
}
|
||||
s.push( Tree( tree.left, middle, ( tree.dimension + 1 ) % k ) );
|
||||
s.push( Tree( middle + 1, tree.right, ( tree.dimension + 1 ) % k ) );
|
||||
}
|
||||
}
|
||||
|
||||
~StaticKDTree(){
|
||||
delete[] kdtree;
|
||||
}
|
||||
~StaticKDTree(){
|
||||
delete[] kdtree;
|
||||
}
|
||||
|
||||
bool NearestNeighbor( InputPoint* result, const InputPoint& point, double radius = std::numeric_limits< T >::max() ) {
|
||||
Metric distance;
|
||||
bool found = false;
|
||||
double nearestDistance = radius;
|
||||
std::stack< NNTree > s;
|
||||
s.push ( NNTree ( 0, size, 0, boundingBox ) );
|
||||
while ( !s.empty() ) {
|
||||
NNTree tree = s.top();
|
||||
s.pop();
|
||||
bool NearestNeighbor( InputPoint* result, const InputPoint& point, double radius = std::numeric_limits< T >::max() ) {
|
||||
Metric distance;
|
||||
bool found = false;
|
||||
double nearestDistance = radius;
|
||||
std::stack< NNTree > s;
|
||||
s.push ( NNTree ( 0, size, 0, boundingBox ) );
|
||||
while ( !s.empty() ) {
|
||||
NNTree tree = s.top();
|
||||
s.pop();
|
||||
|
||||
if ( distance( tree.box, point.coordinates ) >= nearestDistance )
|
||||
continue;
|
||||
if ( distance( tree.box, point.coordinates ) >= nearestDistance )
|
||||
continue;
|
||||
|
||||
if ( tree.left == tree.right )
|
||||
continue;
|
||||
if ( tree.right - tree.left < KDTREE_BASESIZE ) {
|
||||
for ( unsigned i = tree.left; i < tree.right; i++ ) {
|
||||
double newDistance = distance( kdtree[i].coordinates, point.coordinates );
|
||||
if ( newDistance < nearestDistance ) {
|
||||
nearestDistance = newDistance;
|
||||
*result = kdtree[i];
|
||||
found = true;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
Iterator middle = tree.left + ( tree.right - tree.left ) / 2;
|
||||
Iterator middle = tree.left + ( tree.right - tree.left ) / 2;
|
||||
|
||||
double newDistance = distance( kdtree[middle].coordinates, point.coordinates );
|
||||
if ( newDistance < nearestDistance ) {
|
||||
nearestDistance = newDistance;
|
||||
*result = kdtree[middle];
|
||||
found = true;
|
||||
}
|
||||
double newDistance = distance( kdtree[middle].coordinates, point.coordinates );
|
||||
if ( newDistance < nearestDistance ) {
|
||||
nearestDistance = newDistance;
|
||||
*result = kdtree[middle];
|
||||
found = true;
|
||||
}
|
||||
|
||||
Less comperator( tree.dimension );
|
||||
if ( !comperator( point, kdtree[middle] ) ) {
|
||||
NNTree first( middle + 1, tree.right, ( tree.dimension + 1 ) % k, tree.box );
|
||||
NNTree second( tree.left, middle, ( tree.dimension + 1 ) % k, tree.box );
|
||||
first.box.min[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
second.box.max[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
s.push( second );
|
||||
s.push( first );
|
||||
}
|
||||
else {
|
||||
NNTree first( middle + 1, tree.right, ( tree.dimension + 1 ) % k, tree.box );
|
||||
NNTree second( tree.left, middle, ( tree.dimension + 1 ) % k, tree.box );
|
||||
first.box.min[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
second.box.max[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
s.push( first );
|
||||
s.push( second );
|
||||
}
|
||||
}
|
||||
|
||||
return found;
|
||||
}
|
||||
Less comperator( tree.dimension );
|
||||
if ( !comperator( point, kdtree[middle] ) ) {
|
||||
NNTree first( middle + 1, tree.right, ( tree.dimension + 1 ) % k, tree.box );
|
||||
NNTree second( tree.left, middle, ( tree.dimension + 1 ) % k, tree.box );
|
||||
first.box.min[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
second.box.max[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
s.push( second );
|
||||
s.push( first );
|
||||
}
|
||||
else {
|
||||
NNTree first( middle + 1, tree.right, ( tree.dimension + 1 ) % k, tree.box );
|
||||
NNTree second( tree.left, middle, ( tree.dimension + 1 ) % k, tree.box );
|
||||
first.box.min[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
second.box.max[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
|
||||
s.push( first );
|
||||
s.push( second );
|
||||
}
|
||||
}
|
||||
return found;
|
||||
}
|
||||
|
||||
private:
|
||||
typedef unsigned Iterator;
|
||||
struct Tree {
|
||||
Iterator left;
|
||||
Iterator right;
|
||||
unsigned dimension;
|
||||
Tree() {}
|
||||
Tree( Iterator l, Iterator r, unsigned d ): left( l ), right( r ), dimension( d ) {}
|
||||
};
|
||||
struct NNTree {
|
||||
Iterator left;
|
||||
Iterator right;
|
||||
unsigned dimension;
|
||||
BoundingBox< k, T > box;
|
||||
NNTree() {}
|
||||
NNTree( Iterator l, Iterator r, unsigned d, const BoundingBox< k, T >& b ): left( l ), right( r ), dimension( d ), box ( b ) {}
|
||||
};
|
||||
class Less {
|
||||
public:
|
||||
Less( unsigned d ) {
|
||||
dimension = d;
|
||||
assert( dimension < k );
|
||||
}
|
||||
typedef unsigned Iterator;
|
||||
struct Tree {
|
||||
Iterator left;
|
||||
Iterator right;
|
||||
unsigned dimension;
|
||||
Tree() {}
|
||||
Tree( Iterator l, Iterator r, unsigned d ): left( l ), right( r ), dimension( d ) {}
|
||||
};
|
||||
struct NNTree {
|
||||
Iterator left;
|
||||
Iterator right;
|
||||
unsigned dimension;
|
||||
BoundingBox< k, T > box;
|
||||
NNTree() {}
|
||||
NNTree( Iterator l, Iterator r, unsigned d, const BoundingBox< k, T >& b ): left( l ), right( r ), dimension( d ), box ( b ) {}
|
||||
};
|
||||
class Less {
|
||||
public:
|
||||
Less( unsigned d ) {
|
||||
dimension = d;
|
||||
assert( dimension < k );
|
||||
}
|
||||
|
||||
bool operator() ( const InputPoint& left, const InputPoint& right ) {
|
||||
assert( dimension < k );
|
||||
return left.coordinates[dimension] < right.coordinates[dimension];
|
||||
}
|
||||
private:
|
||||
unsigned dimension;
|
||||
};
|
||||
bool operator() ( const InputPoint& left, const InputPoint& right ) {
|
||||
assert( dimension < k );
|
||||
return left.coordinates[dimension] < right.coordinates[dimension];
|
||||
}
|
||||
private:
|
||||
unsigned dimension;
|
||||
};
|
||||
|
||||
BoundingBox< k, T > boundingBox;
|
||||
InputPoint* kdtree;
|
||||
Iterator size;
|
||||
BoundingBox< k, T > boundingBox;
|
||||
InputPoint* kdtree;
|
||||
Iterator size;
|
||||
};
|
||||
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user