Java: KD Tree

This is a Java implementation of a KD-Tree construction.  It's based off of the Wikipedia listing.  If you need some help in understanding what is actually happening, check this out - it shows a live view of the tree's construction and now the graph gets split on the various dimensions.


/**
* Build a KD tree of PointNodes (contains x,y coordinates and Left/Right nodes)
* based on a list of Points (contains x and y coordinates)
 * This is for a 2D tree but could easily be turned into a higher dimension of tree.
*
* @param pointList Random list of points on a 2D plane
* @param depth current depth of the tree
* @return the head of the KD tree
*/
public static PointNode kdtree(List<Point> pointList, int depth)
{
if (pointList == null || pointList.size() == 0)
return null;
//Choose the required dimention (x, y, etc)
int axis = depth % 2;

//Sort all points by their value on the given axis
Collections.sort(pointList, pointCompareMap.get(axis));

int median = pointList.size() / 2;

//use the middle as the root
PointNode root = new PointNode(pointList.get(median), depth);

//Make kdtrees of the left and right sides!
root.setLeft(kdtree(pointList.subList(0, median), depth + 1));
root.setRight(kdtree(pointList.subList(median + 1, pointList.size()),
depth + 1));

return root;
}

Note: the pointCompareMap tells you what dimension to split on based off of your current depth.

//Global
private static Map<Integer, Comparator<Point>> pointCompareMap;



//Defined in Main or the Constructor.
pointCompareMap = new HashMap<Integer, Comparator<Point>>()
{
{
put(0, new PointNodeComparator0());
put(1, new PointNodeComparator1());
//Add more here for a higher dimension (3D, 4D, etc.) of tree
}
};