Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
This script compares the performance of three different methods for finding the nearest neighbors from a list of coordinates. The first method is a simple 'for loop'. The second method uses the KDTree object from the scipy.spatial library. The third method uses cKDTree from the scipy.spatial library. Both the KDTree and cKDTree methods are signi…
import numpy as np
import scipy.spatial as ss
from itertools import combinations
from itertools import permutations
import math
import random
from time import clock
def length(point1, point2):
return math.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)
def find_node_v1(points, nodeCount, node):
best_seg = 10000000
for i in xrange(0, nodeCount):
if node != i:
seg_len = length(points[node], points[i])
if seg_len < best_seg:
best_seg = seg_len
best_nieghbor = i
return best_nieghbor
def get_kdtree(points):
tree = ss.KDTree(points)
return tree
def get_ckdtree(points):
tree = ss.cKDTree(points)
return tree
def query_tree(tree, points, node):
q_result = tree.query(points[node], k=2)
return q_result[1][1]
def main():
nodeCount = 10000
# create a random set of coordinates
points = []
for i in xrange(0, nodeCount):
points.append((random.randint(0, nodeCount), random.randint(0, nodeCount)))
# Run a time trial for three different methods for finding the nearest nieghbor
# for nodes 0-1000
# Method 1 using a for loop
t0 = clock()
nieghbors1 = []
for i in xrange(0,100):
ind = find_node_v1(points, nodeCount, i)
t1 = clock()
print "Method 1 (using for loops) time: ", t1-t0
# Method 2 using KDTree
t0 = clock()
nieghbors2 = []
tree = get_kdtree(points)
for i in xrange(0,100):
ind = query_tree(tree, points, i)
t1 = clock()
print "Method 2 (using kdtree) time: ", t1-t0
# Method 3 using cKDTree
t0 = clock()
nieghbors3 = []
tree = get_ckdtree(points)
for i in xrange(0,100):
ind = query_tree(tree, points, i)
t1 = clock()
print "Method 3 (using ckdtree) time: ", t1-t0
assert nieghbors1 == nieghbors2
assert nieghbors1 == nieghbors3
if __name__ == "__main__":
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment