stefankoegl / kdtree

A Python implementation of a kd-tree
ISC License
365 stars 118 forks source link

search_knn() and search_nn() speedup #25

Closed uozuAho closed 8 years ago

uozuAho commented 9 years ago

It's me again. I've made search_knn() and search_nn() much faster.

The fix uses @cyberjoac's bounded priority queue improvement with a bug fix (squared distance comparison).

The speedup is quite substantial, eg. approx. x300 for the following example on my machine:

from __future__ import print_function
import random
import time

# original kdtree
from orig import kdtree as kdorig
# with knn fix
from uozu2 import kdtree as uozukd

def main():
    nn_timing(3000, 100)

def nn_timing(num_points=1000, num_searches=1000, point_range=1e4):
    points = []
    for i in range(num_points):
        x = random.random()*point_range
        y = random.random()*point_range
        points.append((x, y))

    kdo = kdorig.create(points)
    ukd = uozukd.create(points)

    print('{0} points, {1} searches'.format(num_points, num_searches))
    n_random_nn_searches(kdo, num_searches, 'kd original', point_range)
    n_random_nn_searches(ukd, num_searches, 'knn fixed', point_range)

def n_random_nn_searches(kdtree, n, title, point_range=1e4):
    start = time.time()
    for i in range(n):
        x, y = (random.random()*point_range, random.random()*point_range)
        result = kdtree.search_nn((x, y))
    end = time.time()
    print('{0}: {1} seconds'.format(title, end - start))

if __name__ == '__main__':
    main()

output:

3000 points, 100 searches
kd original: 2.53047800064 seconds
knn fixed: 0.00735902786255 seconds

I've also tried to make all the formatting changes you requested in @cyberjoac's pull request. I'm not too good at git so the history may look a little wonky.