k-Dimensional Trees

Computational geometry is generally concerned with algorithms that solve problems that can be stated in purely geometric terms like: find the smallest polygon which encloses a set of points or construct a data structure to support efficient querying of a set of points. I once interviewed at company that relied on methods from this field. It was during an afternoon one-on-one interview that the k-d tree data structure came up. We, the interviewer and I, were discussing how to find the closest point in a point cloud. At the time, I struggled to find the answer (we were in bonus territory, anyway:)). He explained how one could create a data structure that partitioned the point space into finer and finer cells. And if one keeps track of the relationship between a cell and its constituents, then one can quickly search the space. Essentially, that is what a k-d tree does for point clouds.

k-D dimensional trees are data structures that allow for fast search and querying of point sets. These binary trees have nodes which are k-dimensional points. The tree organizes points by partitioning the point space along each coordinate axis.

Driven to understand this data structure, I coded up a C++ implementation, which you can see on github. Looking at the code: Kdtree class (in kdtree.cpp) represents a k-d tree; main.cpp is essentially a test, that creates a k-d tree with ten randomly generated 3-d points.

{8.40924, 4.11543, 8.10499},
{6.50689, 1.3663, 3.43026},
{9.7614, 9.8382, 0.672512},
{0.113181, 2.22785, 3.46726},
{5.23381, 4.69416, 4.74723},
{9.74655, 0.191659, 1.20641},
{2.8546, 7.32662, 8.51895},
{6.21829, 0.779546, 1.82988},
{8.83612, 8.70544, 2.40537},
{6.50697, 2.70078, 1.93852}


Which when ploted looks like:

So how does one go about making the tree? Well, we first have to determine which among the ten points will be the root. One could chose a point at random, but doing so won't necessarily yield a balanced tree. The algorithm procedes by recursively partitioning the points along each axis, which in our example are the x,y,z - axes. Beginning with the x-axis, the list of points is sorted with respect the the x-axis, where the median point is choosen to be the root (denote green in the figure below) Points with x-values less than root are then grouped into what I call the left list. Likewise, points with x-values greater than root are grouped into the right list. Each list represents new and smaller point clouds, which in turn are partitioned, but this time not along the x-axis, but the y-axis. Specificaly, the median with respect to the y-axis. The root's children are constructed from the lists' medians--the median of the left list becomes the left child, and the median of the right list becomes the right child.

And the process continues, recursively partition left and right lists, until such lists are empty. The figure below summarizes the first few partitions, and highlights the nodes created.

The algorithm loosely described above in detail:

void Kdtree::_makeTree( Node* head, std::list<Kdtree::vecType>& plist, int depth )
{
if( !plist.empty() )
{
int k = plist.front().size();
int axis = depth % k;

std::list<Kdtree::vecType> left_list;
std::list<Kdtree::vecType> right_list;
Kdtree::vecType median = Kdtree::findMedian(axis, plist, left_list, right_list);

Node* left_node = new Node(k);
Node* right_node = new Node(k);

Kdtree::_makeTree( left_node, left_list, depth+1);

Kdtree::_makeTree( right_node, right_list, depth+1);

// Using lambda function here, to define comparison function--parametrized by 'axis'