在3d數組中找到2d元素，這些元素類似於另一個3d數組中的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的歐幾里得）。

PS考慮將query_ball_tree()用於非常大的數據集。 在這種情況下，A和B都必須編入索引。

PS我不確定元素的二維性應該有什么影響; 我給出的示例代碼將它們視為1d，並且至少在使用城市街區指標時是相同的。

Answer 2

從np.allclose的文檔中，我們得到：

如果以下等式是元素為True，則allclose返回True。

絕對值（a - b）<=（atol + rtol * absolute（b））

使用該標准，我們可以使用broadcasting進行矢量化實施，根據所述問題進行定制，如下所示 -

# Setup parameters
rtol,atol = 1e-04, 1e-06

# Use np.allclose criteria to detect true/false across all pairwise elements
mask = np.abs(A[:,None,] - B) <= (atol + rtol * np.abs(B))

# Use the problem context to get final output
out = np.nonzero(mask.all(axis=(2,3)).any(1))[0]