在多维空间中高效地寻找邻域，并基于邻近度计算值之和

Question

我的任务是找出中心元素可变距离内所有元素的总值。元素使用3个维度（我的数据中的列）进行排列。每个元素在给定的三维空间中都有一个唯一的位置（并且有一个唯一的id）。你知道吗

我有一个工作版本，做我想要的，但它是非常缓慢。我正在使用itertuples，使用一个子集dataframe查找每个元组的值，apply(np.I关闭)，并使用.at设置值（请参见下面的代码）。你知道吗

问题不在于代码的功能，而在于可伸缩性。因为我想设置一个可变距离来测量，并且我想为每行计算这个值，所以它最终迭代nrows x ndistance，目前每次迭代需要1.7秒（我的数据有25000行，我估计每尝试一次距离大约需要12小时）。你知道吗

import pandas as pd
import numpy as np

数据结构示例：

df = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19], 
                          'x':[-2,-2,-2,-1,-1,-1,-1,0,0,0,0,0,1,1,1,1,2,2,2], 
                          'y':[2,1,0,2,1,0,-1,2,1,0,-1,-2,1,0,-1,-2,0,-1,-2], 
                          'z':[0,1,2,-1,0,1,2,-2,-1,0,1,2,-2,-1,0,1,-2,-1,0], 
                          'val':[0,0,0,1,0,0,6,3,7,11,0,0,14,18,10,4,20,15,2]})
df.set_index('id', inplace=True)
# The 'val' column can have any non-negative whole number, I've just picked some randomly.

“工作”代码：

n = 0  #Initial distance
while n < 3:  #This part allows me to set my distance range
    df['n{0}'.format(n)] = np.nan  #create a column for the new values
    for row in df.itertuples():
        valsum = df[(df['x'].apply(np.isclose, b=row.x, atol=n)) & 
                    (df['y'].apply(np.isclose, b=row.y, atol=n)) & 
                    (df['z'].apply(np.isclose, b=row.z, atol=n))].val.sum()
        df.at[row.Index, 'n{0}'.format(n)] = valsum
    n += 1

电流/期望输出：

    x   y   z   val n0  n1  n2
id                          
1   -2  2   0   0   0   1   22
2   -2  1   1   0   0   0   25
3   -2  0   2   0   0   6   17
4   -1  2   -1  1   1   11  54
5   -1  1   0   0   0   19  70
6   -1  0   1   0   0   17  57
7   -1  -1  2   6   6   6   31
8   0   2   -2  3   3   25  74
9   0   1   -1  7   7   54  99
10  0   0   0   11  11  46  111
11  0   -1  1   0   0   31  73
12  0   -2  2   0   0   10  33
13  1   1   -2  14  14  62  99
14  1   0   -1  18  18  95  105
15  1   -1  0   10  10  60  107
16  1   -2  1   4   4   16  66
17  2   0   -2  20  20  67  100
18  2   -1  -1  15  15  65  101
19  2   -2  0   2   2   31  80

我知道“n0”列等于“val”列，因为搜索距离是0，但我希望显示我要查找的内容。val列中所有项目的总和为111，这与（x，y，z）=（0,0,0）时相同。这是因为（0,0,0）是本例中数据的中心，因此距离为2将捕获所有元素。我想在5-10的带宽范围内这样做。你知道吗

我的终极问题是：我如何做到这一点，但更快/更有效？你知道吗

Answer 1

在k维空间中寻找最近邻是k-d树数据结构（Wikipedia）的一个经典例子。sciketlearn有一个灵活的实现（docs），我在下面使用它，因为您的问题中使用的条件逻辑似乎定义了Chebyshev距离度量（Wikipedia），sciketlearn本机支持这个度量。SciPy的cKDTree（docs，C++ source code）只支持欧几里德（L2）距离度量，但是针对它进行了优化，因此可能更快。你知道吗

# Setup
df = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19], 
                   'x':[-2,-2,-2,-1,-1,-1,-1,0,0,0,0,0,1,1,1,1,2,2,2], 
                   'y':[2,1,0,2,1,0,-1,2,1,0,-1,-2,1,0,-1,-2,0,-1,-2], 
                   'z':[0,1,2,-1,0,1,2,-2,-1,0,1,2,-2,-1,0,1,-2,-1,0], 
                   'val':[0,0,0,1,0,0,6,3,7,11,0,0,14,18,10,4,20,15,2]})
df.set_index('id', inplace=True)


from sklearn.neighbors import KDTree

# Build k-d tree with the Chebyshev metric, AKA L-infinity
tree = KDTree(df[['x', 'y', 'z']].values, metric='chebyshev')

for radius in [0, 1, 2]:
    # Populate new column with placeholder integer
    df[f'n{radius}'] = -1
    for i, row in df.iterrows():
        coords = row[['x', 'y', 'z']].values.reshape(1, -1)
        idx = tree.query_radius(coords, r=radius)[0]
        df.loc[i, f'n{radius}'] = df.iloc[idx]['val'].sum()

df
    x  y  z  val  n0  n1   n2
id                           
1  -2  2  0    0   0   1   22
2  -2  1  1    0   0   0   25
3  -2  0  2    0   0   6   17
4  -1  2 -1    1   1  11   54
5  -1  1  0    0   0  19   70
6  -1  0  1    0   0  17   57
7  -1 -1  2    6   6   6   31
8   0  2 -2    3   3  25   74
9   0  1 -1    7   7  54   99
10  0  0  0   11  11  46  111
11  0 -1  1    0   0  31   73
12  0 -2  2    0   0  10   33
13  1  1 -2   14  14  62   99
14  1  0 -1   18  18  95  105
15  1 -1  0   10  10  60  107
16  1 -2  1    4   4  16   66
17  2  0 -2   20  20  67  100
18  2 -1 -1   15  15  65  101
19  2 -2  0    2   2  31   80