KDtree : bad performances when the point is far away - c++

Question

I have a basic KDtree implementation so far and the performance is very bad when the points are very far, i.e. if all the points are in [-10;10]^3 and the query is for example at (100, 100, 100), then the bruteforce is 20% faster than the normal version. Why?
This implementation is very basic, I know there is possible improvement everywhere, but just to make work with the minimum vital but still, it shouldn't be that slow... right?
KDTree.h
#pragma once
#include <memory>
#include <vector>
#include <cassert>
#include <vector>
#include <stdexcept>
#include <cmath>
#include <stack>
/**
* #brief K-d tree of 3D float points.
*
* #see Variable names, and algorithm inspired by https://youtu.be/TrqK-atFfWY?t=2567
*/
class KDTree
{
public:
/**** Utility structures ****/
struct Point
{
float x, y, z;
float& operator[](int i)
{
assert(i < 3 && i >= 0);
switch(i) {
case 0: return x;
case 1: return y;
case 2: return z;
}
throw std::runtime_error("Invalid index");
}
const float& operator[](int i) const
{
return const_cast<Point&>(*this)[i];
}
float distance(const Point& other) const
{
return sqrtf(distanceSquared(other));
}
float distanceSquared(const Point& other) const
{
const float dx = (x - other.x);
const float dy = (y - other.y);
const float dz = (z - other.z);
return dx*dx + dy*dy + dz*dz;
}
template<typename T> friend T& operator<<(T& lhs, const Point& rhs)
{
lhs << "(" << rhs.x << ", " << rhs.y << ", " << rhs.z << ")";
return lhs;
}
};
struct AABB
{
Point min, max;
bool inside(const Point& p) const
{
return p.x >= min.x && p.y >= min.y && p.z >= min.z
&& p.x < max.x && p.y < max.y && p.z < max.z;
}
};
public:
/**** Utility (static) functions ****/
/**
* #brief Compute the bounding box of a set of points.
* #param points
* The points to make the boudning box from.
* If the set of points is empty, the returned value is undefined.
* #return
* The bounding box of theses points.
* The bounding box will be always of minimum size, i.e. there will be a vertex on each of the 6 faces of the
* returned AABB.
*/
static AABB computeBoundingBox(const std::vector<Point> &points);
static float median(std::vector<float> vec);
static void store_min(float& current, float newValue)
{
if(current > newValue) {
current = newValue;
}
}
static void store_max(float& current, float newValue)
{
if(current < newValue) {
current = newValue;
}
}
public:
KDTree() = default;
explicit KDTree(const std::vector<Point>& points);
Point computeNearestNeighbor(const Point& pos) const;
private:
/**
* #brief Represent each node of the k-d tree.
*/
struct Node
{
/**
* #brief
* Children. Both of them or neither of them are null.
*/
std::unique_ptr<Node> left, right;
/**
* The distance between the origin and the wall to split (for left node),
* or the distance from the wall and the end to split (for right node).
*/
float splitDistance;
int splitDim;
std::vector<Point> points;
/**
* #return true if this node is a leaf node (it has no children).
*/
bool leaf() const
{
// left or right it doesn't matter
return left == nullptr;
}
};
struct SplitStack
{
int dim;
std::vector<Point> points;
Node* node;
AABB aabb;
};
private:
/**
* #brief Split the tree during the build.
* #param dim The dimension to split.
* #param points The list of remaining candidate points for this area, inside the AABB.
* #param node The node, allocated, to fill.
* #param aabb The bounding box of the node.
*/
void split(std::stack<SplitStack>& stack);
void searchRecursive(const Point& pos, Node* node, float& currentDist, Point& currentNeighbor) const;
private:
AABB m_rootAABB;
std::unique_ptr<Node> m_root;
};
KDTree.cpp
#include "KDTree.h"
#include <climits>
#include <cfloat>
#include <algorithm>
#include <iostream>
#include <cmath>
float KDTree::median(std::vector<float> vec)
{
size_t size = vec.size();
if (size == 0)
{
return 0; // Undefined, really.
}
else
{
std::sort(vec.begin(), vec.end());
if (size % 2 == 0)
{
return (vec[size / 2 - 1] + vec[size / 2]) / 2;
}
else
{
return vec[size / 2];
}
}
}
KDTree::AABB KDTree::computeBoundingBox(const std::vector<Point>& points)
{
AABB res;
if (!points.empty())
{
const auto& firstPoint = points.front();
// Initialize the bounding box to a point
for (int dim = 0; dim < 3; dim++)
{
res.min[dim] = firstPoint[dim];
res.max[dim] = firstPoint[dim];
}
// Grow the bounding box for each point if needed
for (const Point& point: points)
{
for (int dim = 0; dim < 3; dim++)
{
store_min(res.min[dim], point[dim]);
store_max(res.max[dim], point[dim]);
}
}
}
return res;
}
KDTree::KDTree(const std::vector<Point>& points)
{
m_rootAABB = computeBoundingBox(points);
m_root = std::make_unique<Node>();
std::stack<SplitStack> stack;
stack.push(SplitStack{0, points, m_root.get(), m_rootAABB});
split(stack);
}
void KDTree::split(std::stack<SplitStack>& stack)
{
// Split recursively in x, y, z, x, y, z...
// Split at the center
// dim axis -->
// 0 --------- aabb[dim].min --------------------------- aabb[dim].max --------- +inf
// ------------------|--------------------|-------------------|-------------------
// ---------------------- left node ----------- right node -----------------------
// <-----------------> <--------------->
// splitDistance: if left if right
while (!stack.empty())
{
std::vector<Point> points = std::move(stack.top().points);
Node& node = *stack.top().node;
int dim = stack.top().dim;
AABB aabb = stack.top().aabb;
stack.pop();
// Stop condition
if (points.size() > 100)
{
node.splitDim = dim;
// Absolute position in the dimension of the split
node.splitDistance = (aabb.max[dim] + aabb.min[dim]) / 2.0f;
AABB leftAABB = aabb;
leftAABB.max[dim] = node.splitDistance;
AABB rightAABB = aabb;
rightAABB.min[dim] = leftAABB.max[dim];
std::vector<Point> leftPoints, rightPoints;
for (const Point& p: points)
{
if (leftAABB.inside(p))
{
leftPoints.push_back(p);
}
else
{
rightPoints.push_back(p);
}
}
const int nextDim = (dim + 1) % 3;
node.right = std::make_unique<Node>();
stack.push(SplitStack{nextDim, std::move(rightPoints), node.right.get(), rightAABB});
node.left = std::make_unique<Node>();
stack.push(SplitStack{nextDim, std::move(leftPoints), node.left.get(), leftAABB});
}
else
{
// Leaf
node.points = std::move(points);
}
}
}
KDTree::Point KDTree::computeNearestNeighbor(const KDTree::Point& pos) const
{
// Are we left or right?
const Node *node = m_root.get();
AABB aabb = m_rootAABB;
float dist = FLT_MAX;
Point res;
searchRecursive(pos, m_root.get(), dist, res);
return res;
}
void KDTree::searchRecursive(const Point& pos, Node *node, float& currentDist, Point& currentNeighbor) const
{
// Are we on a leaf?
if (node->leaf())
{
// We are on a leaf
// Search brute force into the leaf node
for (const auto& other: node->points)
{
const float d = other.distanceSquared(pos);
if (d < currentDist)
{
currentDist = d;
currentNeighbor = other;
}
}
}
else
{
Node *front, *back;
// Are we on the left side?
if (pos[node->splitDim] < node->splitDistance)
{
// Pos is on the left side
front = node->left.get();
back = node->right.get();
}
else
{
// Pos is on the right side
front = node->right.get();
back = node->left.get();
}
searchRecursive(pos, front, currentDist, currentNeighbor);
// If the current closest point is closer than the closest point of the back face, no need to search in the back
// face because it will be always further.
// If not, we save half of the time for the current node
const float backDist = fabsf(node->splitDistance - pos[node->splitDim]);
// Do not forget all distances all squared
if (backDist * backDist <= currentDist)
{
// If it can be closer, search also in this node
searchRecursive(pos, back, currentDist, currentNeighbor);
}
}
}
main.cpp
#include <iostream>
#include <vector>
#include <random>
#include "KDTree.h"
#include "viewer.h"
#include <chrono>
class Timer {
public:
Timer(const std::string& title) : title_(title), beg_(clock_::now()) {}
~Timer() {
std::cout << title_ << " elapsed: " << elapsed() << "s" << std::endl;
}
void reset() { beg_ = clock_::now(); }
double elapsed() const {
return std::chrono::duration_cast<second_> (clock_::now() - beg_).count();
}
private:
std::string title_;
typedef std::chrono::high_resolution_clock clock_;
typedef std::chrono::duration<double, std::ratio<1> > second_;
std::chrono::time_point<clock_> beg_;
};
std::vector<KDTree::Point> randomPoints(int size, float bounds = 10.0f)
{
std::vector<KDTree::Point> points;
std::uniform_real_distribution<float> dist(-bounds, bounds);
std::mt19937 engine;
for(int i = 0; i < size; i++)
{
KDTree::Point point;
for(int dim = 0; dim < 3; dim++) {
point[dim] = dist(engine);
}
points.push_back(point);
}
return points;
}
int main()
{
auto points = randomPoints(1'000'000);
for(int i = 0; i < points.size(); i++) {
auto& point = points[i];
}
auto aabb = KDTree::computeBoundingBox(points);
KDTree kdtree;
{
Timer timer("KDTree build");
kdtree = KDTree(points);
}
{
const int N = 1'000;
auto testPts = randomPoints(N, 100.0f);
std::vector<KDTree::Point> resKD(N), resBrute(N);
{
Timer timer("KDTree");
for(int i = 0; i < testPts.size(); i++) {
resKD[i] = kdtree.computeNearestNeighbor(testPts[i]);
}
}
{
Timer timer("Bruteforce");
for(int i = 0; i < testPts.size(); i++) {
const auto& test = testPts[i];
// Brute force
KDTree::Point cur;
float curDist = FLT_MAX;
for(const auto& brute : points) {
if(brute.distanceSquared(test) < curDist) {
curDist = brute.distanceSquared(test);
cur = brute;
}
}
resBrute[i] = cur;
}
}
{
float delta = 0.0f;
for(int i = 0; i < N; i++) {
delta += resKD[i].x - resBrute[i].x;
delta += resKD[i].y - resBrute[i].y;
delta += resKD[i].z - resBrute[i].z;
}
std::cout << "delta = " << delta << std::endl;
}
}
return 0;
}
Output:
KDTree build elapsed: 0.190593s
KDTree elapsed: 2.69598s
Bruteforce elapsed: 2.34136s
delta = 0
When I lower to randomPoints(N, 10.0f):
KDTree build elapsed: 0.195519s
KDTree elapsed: 0.000914431s
Bruteforce elapsed: 2.35679s
delta = 0

Answer 1

As written on wikipedia:
Additionally, even in low-dimensional space, if the average pairwise distance between the k nearest neighbors of the query point is significantly less than the average distance between the query point and each of the k nearest neighbors, the performance of nearest neighbor search degrades towards linear, since the distances from the query point to each nearest neighbor are of similar magnitude. (In the worst case, consider a cloud of points distributed on the surface of a sphere centered at the origin. Every point is equidistant from the origin, so a search for the nearest neighbor from the origin would have to iterate through all points on the surface of the sphere to identify the nearest neighbor – which in this case is not even unique.)
So this is normal to see performance decrease in a KD Tree when the points are far.