Ближайшие соседи в заданном диапазоне

Question

Я столкнулся с проблемой быстрого поиска ближайших соседей в заданном диапазоне.

Пример набора данных:

id | string | float
0  |   AA   |  0.1
12 |   BB   |  0.5
2  |   CC   |  0.3
102|   AA   |  1.1
33 |   AA   |  2.8
17 |   AA   |  0.5

Для каждой строки выведите количество строк, удовлетворяющих следующим условиям:

1. строковое поле равно текущему
2. поле с плавающей точкой <= текущее значение с плавающей точкой - del </li>

Для этого примера с del = 1.5:

id | count
0  |  0
12 |  0
2  |  0
102|  2  (string is equal row with id=0,33,17 but only in row id=0,17 float value: 1.1-1.5<=0.1, 1.1-1.5<=0.5)
33 |  0  (string is equal row with id=0,102,17 but 2.8-1.5>=0.1/1.1/1.5)
17 |  1  

Чтобы решить эту проблему, я использовал класс BallTree с пользовательской метрикой, но он работает очень долго из-за обхода дерева по обратному пути (для большого набора данных).Может кто-нибудь предложить другие решения или как вы можете увеличить скорость пользовательских метрик до скорости метрики от sklearn.neighbors.DistanceMetric?

Мой код:

from sklearn.neighbors import BallTree
def distance(x, y):
    if(x[0]==y[0] and x[1]>y[1]):
        return (x[1] - y[1])
    else:
        return (x[1] + y[1])

tree2 = BallTree(X, leaf_size=X.shape[0], metric=distance) 
mas=tree2.query_radius(X, r=del, count_only = True)