04_kdtree.py

from scipy.spatial import cKDTree

def kdtree(data, delta=0.1):
    """
    Constructs a 2D k-d-tree from the input array and queries the points within a square around a given point.
    """
    array, query_points = data
    tree = cKDTree(array)
    count = 0
    for point in query_points:
        idx = tree.query_ball_point((point[0], point[1], point[2]), 0.1, p=np.float('inf'))
        filtered = array[idx]
        count += len(filtered)

    return count

array, query_points = data
tree = cKDTree(array)

# Time the construction of the tree
print('Tree construction:\t', end='')
%timeit cKDTree(array)

# Time the querying time
point = query_points[0]
print('Query time:\t\t', end='')
%timeit tree.query_ball_point((point[0], point[1], point[2]), 0.1, p=np.float('inf'))

# Time the total time for the whole array
print('Total time:\t\t', end='')
%timeit kdtree(data)

print('Count for k-d-tree: {:,}'.format(kdtree(data)))