trees/thirdGenKD/KdTree.java

package ags.utils.dataStructures.trees.thirdGenKD;

import ags.utils.dataStructures.BinaryHeap;
import ags.utils.dataStructures.MaxHeap;
import ags.utils.dataStructures.MinHeap;

/**
 *
 */
public class KdTree<T> extends KdNode<T> {
    public KdTree(int dimensions) {
        this(dimensions, 24);
    }

    public KdTree(int dimensions, int bucketCapacity) {
        super(dimensions, bucketCapacity);
    }

    public NearestNeighborIterator<T> getNearestNeighborIterator(double[] searchPoint, int maxPointsReturned, DistanceFunction distanceFunction) {
        return new NearestNeighborIterator<T>(this, searchPoint, maxPointsReturned, distanceFunction);
    }

    public MaxHeap<T> findNearestNeighbors(double[] searchPoint, int maxPointsReturned, DistanceFunction distanceFunction) {
        BinaryHeap.Min<KdNode<T>> pendingPaths = new BinaryHeap.Min<KdNode<T>>();
        BinaryHeap.Max<T> evaluatedPoints = new BinaryHeap.Max<T>();
        int pointsRemaining = Math.min(maxPointsReturned, size());
        pendingPaths.offer(0, this);

        while (pendingPaths.size() > 0 && (evaluatedPoints.size() < pointsRemaining || (pendingPaths.getMinKey() < evaluatedPoints.getMaxKey()))) {
            nearestNeighborSearchStep(pendingPaths, evaluatedPoints, pointsRemaining, distanceFunction, searchPoint);
        }

        return evaluatedPoints;
    }

    @SuppressWarnings("unchecked")
    protected static <T> void nearestNeighborSearchStep (
            MinHeap<KdNode<T>> pendingPaths, MaxHeap<T> evaluatedPoints, int desiredPoints,
            DistanceFunction distanceFunction, double[] searchPoint) {
        // If there are pending paths possibly closer than the nearest evaluated point, check it out
        KdNode<T> cursor = pendingPaths.getMin();
        pendingPaths.removeMin();

        // Descend the tree, recording paths not taken
        while (!cursor.isLeaf()) {
            KdNode<T> pathNotTaken;
            if (searchPoint[cursor.splitDimension] > cursor.splitValue) {
                pathNotTaken = cursor.left;
                cursor = cursor.right;
            }
            else {
                pathNotTaken = cursor.right;
                cursor = cursor.left;
            }
            double otherDistance = distanceFunction.distanceToRect(searchPoint, pathNotTaken.minBound, pathNotTaken.maxBound);
            // Only add a path if we either need more points or it's closer than furthest point on list so far
            if (evaluatedPoints.size() < desiredPoints || otherDistance <= evaluatedPoints.getMaxKey()) {
                pendingPaths.offer(otherDistance, pathNotTaken);
            }
        }

        if (cursor.singlePoint) {
            double nodeDistance = distanceFunction.distance(cursor.points[0], searchPoint);
            // Only add a point if either need more points or it's closer than furthest on list so far
            if (evaluatedPoints.size() < desiredPoints || nodeDistance <= evaluatedPoints.getMaxKey()) {
                for (int i = 0; i < cursor.size(); i++) {
                    T value = (T) cursor.data[i];

                    // If we don't need any more, replace max
                    if (evaluatedPoints.size() == desiredPoints) {
                        evaluatedPoints.replaceMax(nodeDistance, value);
                    } else {
                        evaluatedPoints.offer(nodeDistance, value);
                    }
                }
            }
        } else {
            // Add the points at the cursor
            for (int i = 0; i < cursor.size(); i++) {
                double[] point = cursor.points[i];
                T value = (T) cursor.data[i];
                double distance = distanceFunction.distance(point, searchPoint);
                // Only add a point if either need more points or it's closer than furthest on list so far
                if (evaluatedPoints.size() < desiredPoints) {
                    evaluatedPoints.offer(distance, value);
                } else if (distance < evaluatedPoints.getMaxKey()) {
                    evaluatedPoints.replaceMax(distance, value);
                }
            }
        }
    }
}
1	torben	2744	package ags.utils.dataStructures.trees.thirdGenKD;
2
3			import ags.utils.dataStructures.BinaryHeap;
4			import ags.utils.dataStructures.MaxHeap;
5			import ags.utils.dataStructures.MinHeap;
6
7			/**
8			*
9			*/
10			public class KdTree<T> extends KdNode<T> {
11			public KdTree(int dimensions) {
12			this(dimensions, 24);
13			}
14
15			public KdTree(int dimensions, int bucketCapacity) {
16			super(dimensions, bucketCapacity);
17			}
18
19			public NearestNeighborIterator<T> getNearestNeighborIterator(double[] searchPoint, int maxPointsReturned, DistanceFunction distanceFunction) {
20			return new NearestNeighborIterator<T>(this, searchPoint, maxPointsReturned, distanceFunction);
21			}
22
23			public MaxHeap<T> findNearestNeighbors(double[] searchPoint, int maxPointsReturned, DistanceFunction distanceFunction) {
24			BinaryHeap.Min<KdNode<T>> pendingPaths = new BinaryHeap.Min<KdNode<T>>();
25			BinaryHeap.Max<T> evaluatedPoints = new BinaryHeap.Max<T>();
26			int pointsRemaining = Math.min(maxPointsReturned, size());
27			pendingPaths.offer(0, this);
28
29			while (pendingPaths.size() > 0 && (evaluatedPoints.size() < pointsRemaining \|\| (pendingPaths.getMinKey() < evaluatedPoints.getMaxKey()))) {
30			nearestNeighborSearchStep(pendingPaths, evaluatedPoints, pointsRemaining, distanceFunction, searchPoint);
31			}
32
33			return evaluatedPoints;
34			}
35
36			@SuppressWarnings("unchecked")
37			protected static <T> void nearestNeighborSearchStep (
38			MinHeap<KdNode<T>> pendingPaths, MaxHeap<T> evaluatedPoints, int desiredPoints,
39			DistanceFunction distanceFunction, double[] searchPoint) {
40			// If there are pending paths possibly closer than the nearest evaluated point, check it out
41			KdNode<T> cursor = pendingPaths.getMin();
42			pendingPaths.removeMin();
43
44			// Descend the tree, recording paths not taken
45			while (!cursor.isLeaf()) {
46			KdNode<T> pathNotTaken;
47			if (searchPoint[cursor.splitDimension] > cursor.splitValue) {
48			pathNotTaken = cursor.left;
49			cursor = cursor.right;
50			}
51			else {
52			pathNotTaken = cursor.right;
53			cursor = cursor.left;
54			}
55			double otherDistance = distanceFunction.distanceToRect(searchPoint, pathNotTaken.minBound, pathNotTaken.maxBound);
56			// Only add a path if we either need more points or it's closer than furthest point on list so far
57			if (evaluatedPoints.size() < desiredPoints \|\| otherDistance <= evaluatedPoints.getMaxKey()) {
58			pendingPaths.offer(otherDistance, pathNotTaken);
59			}
60			}
61
62			if (cursor.singlePoint) {
63			double nodeDistance = distanceFunction.distance(cursor.points[0], searchPoint);
64			// Only add a point if either need more points or it's closer than furthest on list so far
65			if (evaluatedPoints.size() < desiredPoints \|\| nodeDistance <= evaluatedPoints.getMaxKey()) {
66			for (int i = 0; i < cursor.size(); i++) {
67			T value = (T) cursor.data[i];
68
69			// If we don't need any more, replace max
70			if (evaluatedPoints.size() == desiredPoints) {
71			evaluatedPoints.replaceMax(nodeDistance, value);
72			} else {
73			evaluatedPoints.offer(nodeDistance, value);
74			}
75			}
76			}
77			} else {
78			// Add the points at the cursor
79			for (int i = 0; i < cursor.size(); i++) {
80			double[] point = cursor.points[i];
81			T value = (T) cursor.data[i];
82			double distance = distanceFunction.distance(point, searchPoint);
83			// Only add a point if either need more points or it's closer than furthest on list so far
84			if (evaluatedPoints.size() < desiredPoints) {
85			evaluatedPoints.offer(distance, value);
86			} else if (distance < evaluatedPoints.getMaxKey()) {
87			evaluatedPoints.replaceMax(distance, value);
88			}
89			}
90			}
91			}
92			}