k-d tree

k-d tree

Given a list of n points, the following algorithm will construct a balanced kd-tree containing those points.

function kdtree (list of points pointList, int depth)
{
    if pointList is empty
        return nil;
    else
    {
        // Select axis based on depth so that axis cycles through all valid values
        var int axis := depth mod k;

        // Sort point list and choose median as pivot element
        select median from pointList;

        // Create node and construct subtrees
        var tree_node node;
        node.location := median;
        node.leftChild := kdtree(points in pointList before median, depth+1);
        node.rightChild := kdtree(points in pointList after median, depth+1);
        return node;
    }
}


This algorithm implemented in the Python programming language is as follows:

class Node:pass
 
def kdtree(pointList, depth=0):
    if not pointList:
        return
 
    # Select axis based on depth so that axis cycles through all valid values
    k = len(pointList[0]) # assumes all points have the same dimension
    axis = depth % k
 
    # Sort point list and choose median as pivot element
    pointList.sort(cmp=lambda x,y:cmp(x[axis],y[axis]))
    median = len(pointList)/2 # choose median
 
    # Create node and construct subtrees
    node = Node()
    node.location = pointList[median]
    node.leftChild = kdtree(pointList[0:median], depth+1)
    node.rightChild = kdtree(pointList[median+1:], depth+1)
    return node
example:

pointList = [(2,3), (5,4), (9,6), (4,7), (8,1), (7,2)]
tree = kdtree(pointList)
Balancing a kd-tree: http://en.wikipedia.org/wiki/Kd-tree