Nearest Neighbor Search: Python

问题:

I have a 2 dimensional array:

MyArray = array([6588252.24, 1933573.3, 212.79, 0, 0],
                [6588253.79, 1933602.89, 212.66, 0, 0],
                 etc...)

The first two elements MyArray[0] and MyArray[1] are the X and Y coordinates of the points.

For every element in the array, I would like to find the quickest way to return its single nearest neighbor in a radius of X units. We are assuming this is in 2D space.

lets say for this example X = 6.

I have solved the problem by comparing every element to every other element, but this takes 15 minutes or so when your list is 22k points long. We hope to eventually run this on lists of about 30million points.

I have read about K-d trees and understand the basic concept, but have had trouble understanding how to script them.

回答1:

Thanks to John Vinyard for suggesting scipy. After some good research and testing, here is the solution to this question:

Prerequisites: Install Numpy and SciPy

1. Import the SciPy and Numpy Modules

2. Make a copy of the 5 dimensional array including just the X and Y values.

3. Create an instance of a cKDTree as such:

   YourTreeName = scipy.spatial.cKDTree(YourArray, leafsize=100)
   #Play with the leafsize to get the fastest result for your dataset
   

4. Query the cKDTree for the Nearest Neighbor within 6 units as such:

   for item in YourArray:
       TheResult = YourTreeName.query(item, k=1, distance_upper_bound=6)
   

   for each item in YourArray, TheResult will be a tuple of the distance between the two points, and the index of the location of the point in YourArray.