python - 在 3d 数组中找到与另一个 3d 数组中的 2d 元素相似的 2d 元素

Question

我有两个 3D 数组，我想识别一个数组中的 2D 元素，这些元素在另一个数组中有一个或多个相似的对应项。

这适用于 Python 3:

import numpy as np
import random

np.random.seed(123)
A = np.round(np.random.rand(25000,2,2),2)
B = np.round(np.random.rand(25000,2,2),2)

a_index = np.zeros(A.shape[0])

for a in range(A.shape[0]):
    for b in range(B.shape[0]):
        if np.allclose(A[a,:,:].reshape(-1, A.shape[1]), B[b,:,:].reshape(-1, B.shape[1]),
                       rtol=1e-04, atol=1e-06):
            a_index[a] = 1
            break

np.nonzero(a_index)[0]

当然，这种方法非常慢。请告诉我，有一种更有效的方法(以及它是什么)。谢谢。

Answer 1

您正在尝试进行全最近邻类型查询。这是具有特殊 O(n log n) 算法的东西，我不知道 python 实现。但是，您可以使用常规的最近邻算法，它也是 O(n log n) 只是慢一点。例如 scipy.spatial.KDTree 或 cKDTree。

import numpy as np
import random
np.random.seed(123)
A = np.round(np.random.rand(25000,2,2),2)
B = np.round(np.random.rand(25000,2,2),2)

import scipy.spatial
tree = scipy.spatial.cKDTree(A.reshape(25000, 4))
results = tree.query_ball_point(B.reshape(25000, 4), r=1e-04, p=1)

print [r for r in results if r != []]
# [[14252], [1972], [7108], [13369], [23171]]

query_ball_point() 不完全等同于 allclose() 但它足够接近，特别是如果您不关心rtol 参数到 allclose()。您还可以选择度量标准(p=1 用于城市街区，或 p=2 用于欧几里德)。

附言考虑对非常大的数据集使用 query_ball_tree()。在这种情况下，A 和 B 都必须被索引。

附言我不确定元素的二维性应该有什么影响；我给出的示例代码将它们视为 1d，至少在使用城市街区指标时是相同的。