Оптимизация поиска ближайшего соседа scipy

Я пытаюсь найти всех ближайших соседей, которые находятся в радиусе 1 км. Вот мой скрипт для построения дерева и поиска ближайших точек,

from pysal.cg.kdtree import KDTree

def construct_tree(s):
    data_geopoints = [tuple(x) for x in s[['longitude','latitude']].to_records(index=False)]
    tree = KDTree(data_geopoints, distance_metric='Arc', radius=pysal.cg.RADIUS_EARTH_KM)
    return tree

def get_neighbors(s,tree):
    indices = tree.query_ball_point(s, 1)
    return indices

#Constructing the tree for search
tree = construct_tree(data)

#Finding the nearest neighbours within 1KM
data['neighborhood'] = data['lat_long'].apply(lambda row: get_neighbors(row,tree))

из того, что я прочитал на странице писала, Он говорит -

KD-дерево, построенное поверх функциональности KD-дерева в scipy. Если вы используете составляющей 0.12 или больше использует scipy.пространственный.cKDTree, в противном случае использует scipy.пространственный.KDTree.

в моем случае он должен использовать cKDTree. Это работает отлично для образца набора данных, но с tree.query_ball_point возвращает список индексов в результате. Каждый список будет иметь 100s элементов. Для моих точек данных (2 миллиона записей) это становится все больше и больше и останавливается из-за проблемы с памятью после определенной точки. Есть идеи, как это решить?