User:Rednaxela/kD-Tree

A nice efficent small kD-Tree. It's quite fast... Feel free to use

/**
 * 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.newtree2;

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;
        ResultHeap resultHeap = new ResultHeap(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);
                    resultHeap.addValue(dist, cursor.locations[i]);
                }
                range = resultHeap.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 = resultHeap.getData();
        double[] dist = resultHeap.getDistances();
        for (int i=0; i<resultHeap.values; 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 ResultHeap {
        private final Object[] data;
        private final double[] distance;
        private final int size;
        private int values;

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

        public void addValue(double dist, Object value) {
            if (values == size && dist >= distance[0]) {
                return;
            }

            // Insert value
            data[values] = value;
            distance[values] = dist;
            values++;

            // Up-Heapify
            for (int c = values-1, p = (c-1)/2; c != 0 && distance[c] > distance[p]; c = p, p = (c-1)/2) {
                Object pData = data[p];
                double pDist = distance[p];
                data[p] = data[c];
                distance[p] = distance[c];
                data[c] = pData;
                distance[c] = pDist;
            }

            // If too big, remove the highest value
            if (values > size) {
                // Move the last entry to the top
                values--;
                data[0] = data[values];
                distance[0] = distance[values];

                // Down-Heapify
                for (int p = 0, c = 1; c < values; p = c,c = p*2+1) {
                    if (c+1 < values && distance[c] < distance[c+1]) {
                        c++;
                    }
                    if (distance[p] < distance[c]) {
                        // Swap the points
                        Object pData = data[p];
                        double pDist = distance[p];
                        data[p] = data[c];
                        distance[p] = distance[c];
                        data[c] = pData;
                        distance[c] = pDist;
                    }
                    else {
                        break;
                    }
                }
            }
        }

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

        public Object[] getData() {
            return data;
        }

        public double[] getDistances() {
            return distance;
        }
    }
}