python - 高效查找具有割断的邻居并返回索引

Question

我在x,y平面上有很多点，长度约为10000，每个点(x,y)都有一个内在半径r.这个小数据集只是我整个数据集的一个小角落。我有一个感兴趣的点(x1,y1)，我想找到(x1,y1)周围1以内的附近点，并且满足之间距离的条件(x,y) 和 (x1,y1) 小于 r。我想返回那些优点的索引，而不是优点本身。

import numpy as np
np.random.seed(2000)
x = 20.*np.random.rand(10000)
y = 20.*np.random.rand(10000)
r = 0.3*np.random.rand(10000)
x1 = 10.  ### (x1,y1) is an interest point 
y1 = 12.
def index_finder(x,y,r,x1,y1):
    idx = (abs(x - x1) < 1.) & (abs(y - y1) < 1.)  ### This cut will probably cut 90% of the data
    x_temp = x[idx]   ### but if I do like this, then I lose the track of the original index
    y_temp = y[idx]
    dis_square = (x_temp - x1)*(x_temp - x1) + (y_temp - y1)*(y_temp - y1)
    idx1 = dis_square < r*r    ### after this cut, there are only a few left 
    x_good = x_temp[idx1]
    y_good = y_temp[idx1]

在这个函数中，我可以找到(x1,y1)周围的优点，但找不到这些优点的索引。但是，我需要原始索引，因为原始索引用于提取与坐标 (x,y) 关联的其他数据。正如我所提到的，样本数据集只是我整个数据集的一小部分，我将在整个数据集上调用上述函数大约 1,000,000 次，因此上述 index_finder 函数的效率为也是一个考虑因素。

对这样的任务有什么想法吗？

Answer 1

方法#1

我们可以简单地用第一个掩码索引到它自己的掩码，以便从第二阶段选择 True 位置掩码值，就像这样 -

idx[idx] = idx1

因此，idx 将具有与原始数组 x 和 y 相对应的最终有效掩码值/有值(value)的位置，即 -

x_good = x[idx]
y_good = y[idx]

然后可以使用该掩码对问题中提到的其他数组进行索引。

<小时/>

方法#2

作为另一种方法，我们可以使用两个条件语句，从而用它们创建两个掩码。最后，将它们与AND-ing组合以获得组合掩码，该掩码可以索引到x和y数组中以获得最终输出。我们不需要以这种方式获取实际索引，因此这是它的另一个好处。

因此，实现 -

X = x-x1
Y = y-y1
mask1 = (np.abs(X) < 1.) & (np.abs(Y) < 1.)
mask2 = X**2 + Y*2 < r**2
comb_mask = mask1 & mask2

x_good = x[comb_mask]
y_good = y[comb_mask]

如果由于某种原因，您仍然需要相应的索引，只需这样做 -

comb_idx = np.flatnonzero(comb_mask)

如果您对同一 x 和 y 数据集的不同 x1 和 y1 对执行这些操作，我建议使用broadcasting通过所有这些x1、y1配对数据集对其进行矢量化，如 this post 中所示。 .