python tips: usage of scipy.spatial.KDTree

kd-tree is a well-known algorithm for searching spatially distributed points. Below code shows how to use kd-tree implemented in scipy.spatial.

#!/usr/bin/env python                                                                                                                                                                                          
#coding:utf-8                                                                                                                                                                                                  
                                                                                                                                                                                                               
import numpy as np                                                                                                                                                                                             
import scipy.spatial as ss                                                                                                                                                                                     
from itertools import combinations
from itertools import permutations

def main():
    x, y = np.mgrid[0:100, 0:100]
    points = zip(x.ravel(), y.ravel())
    tree = ss.KDTree(points)
    # get index of each point which distance from (x=0, y=0) is under 1
    a = tree.query_ball_point([0, 0], 1)
    print [points[i] for i in a]

if __name__ == "__main__":
    main()

output of above code is like:

~$ python above_code.py 
[(0, 0), (0, 1), (1, 0)]