向量化嵌套循环

Question

我希望对嵌套循环进行矢量化处理，该循环将适用于 300,000 个列表的列表，每个列表都包含 3 个值。嵌套循环将每个列表的值与其他列表中的对应值进行比较，并且只会附加对应值最大差异为 0.1 的列表索引。因此，包含 [0.234, 0.456, 0.567] 的列表和包含 [0.246, 0.479, 0.580] 的列表将属于此类别，因为它们的对应值（即 0.234 和 0.246；0.456 和 0.479；0.567 和 0.580）存在差异它们之间小于 0.1。

我目前使用以下嵌套循环来执行此操作，但目前大约需要 58 小时才能完成（总共 90 万亿次迭代）；

import numpy as np
variable = np.random.random((300000,3)).tolist()
out1=list()
out2=list()
for i in range(0:300000):
    for j in range(0:300000):
        if ((i<j) and ((abs(variable[i][0]-variable[j][0]))<0.1) and ((abs(variable[i][1]-variable[j] [1]))<0.1) and ((abs(variable[i][2]-variable[j][2]))<0.1)):
        out1.append(i)  
        out2.append(j)

Answer 1

查看 scipy.spatial;它有很多功能可以有效地解决此类空间查询；尤其是KDTrees，即：

import scipy.spatial
out = scipy.spatial.cKDTree(variable).query_pairs(r=0.1, p=np.infinity)

Answer 2

转换为 NumPy 数组，以便之后更容易使用 NumPy 函数。然后，可以提出两种方法。

方法#1

NumPy 广播可用于将这些扩展到 3D 数组并以矢量化方式执行操作。因此，我们会有这样的实现 -

th = 0.1 # Threshold
arr = np.asarray(variable)
out1,out2 = np.where(np.triu((np.abs(arr[:,None,:] - arr) < th).all(-1),1))

方法 #2

专注于内存效率的替代实现，使用负责此类迭代的选择性索引 -

th = 0.1 # Threshold
arr = np.asarray(variable)
R,C = np.triu_indices(arr.shape[0],1)
mask = (np.abs(arr[R] - arr[C])<th).all(-1)
out1,out2 = R[mask], C[mask]