python - 加快阵列中所有可能对之间的距离

Question

我有几个 (~10^10) 个点的 x、y、z 坐标数组(这里只显示 5 个)

a= [[ 34.45  14.13   2.17]
    [ 32.38  24.43  23.12]
    [ 33.19   3.28  39.02]
    [ 36.34  27.17  31.61]
    [ 37.81  29.17  29.94]]

我想创建一个新数组，其中仅包含与列表中所有其他点至少有一定距离 d 的点。我使用 while 循环写了一段代码，

 import numpy as np
 from scipy.spatial import distance 

 d=0.1 #or some distance 
 i=0
 selected_points=[]
 while i < len(a):
          interdist=[]  
          j=i+1
          while j<len(a):
              interdist.append(distance.euclidean(a[i],a[j]))
              j+=1

          if all(dis >= d for dis in interdist):
              np.array(selected_points.append(a[i]))
          i+=1

这可行，但执行此计算需要很长时间。我在某处读到 while 循环非常慢。

我想知道是否有人对如何加快此计算有任何建议。

编辑:虽然我找到与所有其他粒子至少有一定距离的粒子的目标保持不变，但我只是意识到我的代码中存在严重缺陷，假设我有 3 个粒子，我的代码执行以下操作，对于 i 的第一次迭代，它计算距离 1->2、1->3，假设 1->2 小于阈值距离 d，因此代码会丢弃粒子 1。对于 i 的下一次迭代，它只执行 2->3，假设它发现它大于 d，所以它保留粒子 2，但这是错误的!因为 2 也应该与粒子 1 一起被丢弃。 @svohara 的解决方案是正确的!

Answer 1

对于大数据集和低维点(例如 3 维数据)，有时使用空间索引方法会有很大好处。低维数据的一种流行选择是 k-d 树。

策略是索引数据集。然后使用相同的数据集查询索引，返回每个点的 2 个最近邻。第一个最近的邻居总是点本身(dist=0)，所以我们真的想知道下一个最近的点有多远(第二个最近的邻居)。对于 2-NN > 阈值的那些点，您得到了结果。

from scipy.spatial import cKDTree as KDTree
import numpy as np

#a is the big data as numpy array N rows by 3 cols
a = np.random.randn(10**8, 3).astype('float32')

# This will create the index, prepare to wait...
# NOTE: took 7 minutes on my mac laptop with 10^8 rand 3-d numbers
#  there are some parameters that could be tweaked for faster indexing,
#  and there are implementations (not in scipy) that can construct
#  the kd-tree using parallel computing strategies (GPUs, e.g.)
k = KDTree(a)

#ask for the 2-nearest neighbors by querying the index with the
# same points
(dists, idxs) = k.query(a, 2)
# (dists, idxs) = k.query(a, 2, n_jobs=4)  # to use more CPUs on query...

#Note: 9 minutes for query on my laptop, 2 minutes with n_jobs=6
# So less than 10 minutes total for 10^8 points.

# If the second NN is > thresh distance, then there is no other point
# in the data set closer.
thresh_d = 0.1   #some threshold, equiv to 'd' in O.P.'s code
d_slice = dists[:, 1]  #distances to second NN for each point
res = np.flatnonzero( d_slice >= thresh_d )