User:Rednaxela/kD-Tree

A nice efficent 359-line kD-Tree. It's quite fast...

/**
 * Copyright 2009 Rednaxela
 * 
 * This software is provided 'as-is', without any express or implied
 * warranty. In no event will the authors be held liable for any damages
 * arising from the use of this software.
 * 
 * Permission is granted to anyone to use this software for any purpose,
 * including commercial applications, and to alter it and redistribute it
 * freely, subject to the following restrictions:
 * 
 *    1. The origin of this software must not be misrepresented; you must not
 *    claim that you wrote the original software. If you use this software
 *    in a product, an acknowledgment in the product documentation would be
 *    appreciated but is not required.
 * 
 *    2. This notice may not be removed or altered from any source
 *    distribution.
 */

package ags.utils.newtree;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * An efficent well-optimized kd-tree
 * 
 * @author Rednaxela
 */
public class KdTree<T> {
	// Static variables
	private static final int bucketSize = 32;

	// All types
	private final int dimensions;

	// Root only
	private final HashMap<Object, T> map;
	private double[] weights;

	// Leaf only
	private double[][] locations;
	private int locationCount;

	// Stem only
	private KdTree<T> left, right;
	private int splitDimension;
	private double splitValue;

	// Bounds
	private double[] minLimit, maxLimit;

	/**
	 * Extends the bounds of this node do include a new location
	 */
	private final void extendBounds(double[] location) {
		if (minLimit == null) {
			minLimit = Arrays.copyOf(location, dimensions);
			maxLimit = Arrays.copyOf(location, dimensions);
			return;
		}

		for (int i=0; i<dimensions; i++) {
			if (minLimit[i] > location[i]) {
				minLimit[i] = location[i];
			}
			if (maxLimit[i] < location[i]) {
				maxLimit[i] = location[i];
			}
		}
	}

	/**
	 * Find the widest axis of the bounds of this node
	 */
	private final int findWidestAxis() {
		int widest = 0;
		double width = (maxLimit[0] - minLimit[0]);
		for (int i = 1; i < dimensions; i++) {
			double nwidth = maxLimit[i] - minLimit[i];
			if (nwidth > width) {
				widest = i;
				width = nwidth;
			}
		}
		return widest;
	}

	// Main constructor
	public KdTree(int dimensions) {
		this.dimensions = dimensions;

		// Init as leaf
		this.locations = new double[bucketSize][];
		this.locationCount = 0;

		// Init as root
		this.map = new HashMap<Object, T>();
		this.weights = new double[dimensions];
		Arrays.fill(this.weights, 1.0);
	}

	// Child constructor
	private KdTree(KdTree<T> parent, boolean right) {
		this.dimensions = parent.dimensions;

		// Init as leaf
		this.locations = new double[bucketSize][];
		this.locationCount = 0;

		// Init as non-root
		this.map = null;
	}

	/**
	 * Add a point and associated value to the tree
	 */
	public static <T> void addPoint(KdTree<T> tree, double[] location, T value) {
		KdTree<T> cursor = tree;

		while (cursor.locations == null || cursor.locationCount >= cursor.locations.length) {
			if (cursor.locations != null) {
				cursor.splitDimension = cursor.findWidestAxis();
				cursor.splitValue = (cursor.minLimit[cursor.splitDimension] + cursor.maxLimit[cursor.splitDimension]) * 0.5;

				// Don't split node if it has no width in any axis. Double the bucket size instead
				if ((cursor.minLimit[cursor.splitDimension] - cursor.maxLimit[cursor.splitDimension]) == 0) {
					cursor.locations = Arrays.copyOf(cursor.locations, cursor.locations.length * 2);
					break;
				}

				// Create child leaves
				KdTree<T> left = new KdTree<T>(cursor, false);
				KdTree<T> right = new KdTree<T>(cursor, true);

				// Move locations into children
				for (double[] oldLocation : cursor.locations) {
					if (oldLocation[cursor.splitDimension] > cursor.splitValue) {
						// Right
						right.locations[right.locationCount] = oldLocation;
						right.locationCount++;
						right.extendBounds(oldLocation);
					}
					else {
						// Left
						left.locations[left.locationCount] = oldLocation;
						left.locationCount++;
						left.extendBounds(oldLocation);
					}
				}

				// Make into stem
				cursor.left = left;
				cursor.right = right;
				cursor.locations = null;
			}

			cursor.extendBounds(location);

			if (location[cursor.splitDimension] > cursor.splitValue) {
				cursor = cursor.right;
			}
			else {
				cursor = cursor.left;
			}
		}

		cursor.locations[cursor.locationCount] = location;
		cursor.locationCount++;
		cursor.extendBounds(location);

		tree.map.put(location, value);
	}

	/**
	 * Enumeration representing the status of a node during the running 
	 */
	private static enum Status {
		NONE,
		LEFTVISITED,
		RIGHTVISITED,
		ALLVISITED
	}

	/**
	 * Stores a distance and value to output
	 */
	public static class Entry<T> {
		public final double distance;
		public final T value;
		private Entry(double distance, T value) {
			this.distance = distance;
			this.value = value;
		}
	}

	/**
	 * Calculates the nearest 'count' points to 'location', with an arbitrary weighting on dimensions
	 */
	public static <T> List<Entry<T>> nearestNeighbor(KdTree<T> tree, double[] location, int count, double[] weights) {
		tree.weights = weights;
		return nearestNeighbor(tree, location, count);
	}

	/**
	 * Calculates the nearest 'count' points to 'location'
	 */
	public static <T> List<Entry<T>> nearestNeighbor(KdTree<T> tree, double[] location, int count) {
		KdTree<T> cursor = tree;
		Status status = Status.NONE;
		Stack<KdTree<T>> stack = new Stack<KdTree<T>>();
		Stack<Status> statusStack = new Stack<Status>();
		double range = Double.POSITIVE_INFINITY;
		TopArray top = new TopArray(count); 

		do {
			if (status == Status.ALLVISITED) {
				// At a fully visited part. Move up the tree
				cursor = stack.pop();
				status = statusStack.pop();
				continue;
			}

			if (status == Status.NONE && cursor.locations != null) {
				// At a leaf. Use the data.
				for (int i=0; i<cursor.locationCount; i++) {
					double dist = sqrPointDist(cursor.locations[i], location, tree.weights);
					top.addValue(dist, cursor.locations[i]);
				}
				range = top.getMaxDist();

				if (stack.empty()) {
					break;
				}
				cursor = stack.pop();
				status = statusStack.pop();
				continue;
			}

			// Going to descend
			KdTree<T> nextCursor = null;
			if (status == Status.NONE) {
				// At a fresh node, descend the most probably useful direction
				if (location[cursor.splitDimension] > cursor.splitValue) {
					// Descend right
					nextCursor = cursor.right;
					status = Status.RIGHTVISITED;
				}
				else {
					// Descend left;
					nextCursor = cursor.left;
					status = Status.LEFTVISITED;
				}
			}
			else if (status == Status.LEFTVISITED) {
				// Left node visited, descend right.
				nextCursor = cursor.right;
				status = Status.ALLVISITED;
			}
			else if (status == Status.RIGHTVISITED) {
				// Right node visited, descend left.
				nextCursor = cursor.left;
				status = Status.ALLVISITED;
			}

			// Check if it's worth descending. Assume it is if it's sibling has not been visited yet. 
			if (status == Status.ALLVISITED) {
				if (nextCursor.locationCount == 0 || sqrPointRegionDist(location, nextCursor.minLimit, nextCursor.maxLimit, tree.weights) > range) {
					continue;
				}
			}

			// Descend down the tree
			stack.push(cursor);
			statusStack.push(status);
			cursor = nextCursor;
			status = Status.NONE;
		} while (stack.size() > 0 || status != Status.ALLVISITED);

		ArrayList<Entry<T>> results = new ArrayList<Entry<T>>(count);
		Object[] data = top.getData();
		double[] dist = top.getDistances();
		for (int i=0; i<data.length; i++) {
			T value = tree.map.get(data[i]);
			results.add(new Entry<T>(dist[i], value));
		}

		return results;
	}

	/**
	 * Calculates the (squared euclidean) distance between two points
	 */
	private static final double sqrPointDist(double[] p1, double[] p2, double[] weights) {
		double d = 0;

		for (int i=0; i<p1.length; i++) {
			double diff = (p1[i] - p2[i]) * weights[i];
			d += diff * diff;
		}

		return d;
	}

	/**
	 * Calculates the closest (squared euclidean) distance between in a point and a bounding region
	 */
	private static final double sqrPointRegionDist(double[] point, double[] min, double[] max, double[] weights) {
		double d = 0;

		for (int i=0; i<point.length; i++) {
			if (point[i] > max[i]) {
				double diff = (point[i] - max[i]) * weights[i];
				d += diff * diff;
			} else if (point[i] < min[i]) {
				double diff = (point[i] - min[i]) * weights[i];
				d += diff * diff;
			}
		}

		return d;
	}

	/**
	 * Class for tracking up to 'size' closest values
	 */
	private static class TopArray {
		private final Object[] data;
		private final double[] distance;
		private final int size;
		private int values;

		public TopArray(int size) {
			this.data = new Object[size];
			this.distance = new double[size];
			this.size = size;
			this.values = 0;
		}

		public void addValue(double dist, Object value) {
			int i = values;
			while (i > 0 && distance[i-1] > dist) {
				i--;
			}

			if (i >= size) {
				return;
			}

			if (values < size) {
				values++;
			}

			System.arraycopy(distance, i, distance, i+1, values-(i+1));
			distance[i] = dist;
			System.arraycopy(data, i, data, i+1, values-(i+1));
			data[i] = value;
		}

		public double getMaxDist() {
			if (values < size) {
				return Double.POSITIVE_INFINITY;
			}
			return distance[size-1];
		}

		public Object[] getData() {
			return Arrays.copyOf(data, values);
		}

		public double[] getDistances() {
			return Arrays.copyOf(distance, values);
		}
	}
}