python - 如何有效地找到与第二个数组值匹配的第一个数组值的索引？

Question

我有两个 numpy 数组 A 和 B。A 的形状为 (10000000, 3)，B 的形状为 (1000000, 3)。这两个数组都是 XYZ 坐标，因此 B 对应于 A 的某个区域。我必须找到对应于值 B 的 A 的索引。现在我正在解决如下问题。我需要一些帮助来使用 Numpy 或其他 python 包对其进行优化。

extract_BinA=np.empty(B.shape[0])
for i in range(B.shape[0]):
    for j in range(A.shape[0]):
        if(A[j][0]==B[i][0] and A[j][1]==B[i][1] and A[j][2]==B[i][2]):
            extract_BinA[i]=j

Answer 1

这里的问题不是纯 Python 代码的速度，而是算法 本身。您可以使用sorted-arrays 或hash-tables 将算法的复杂度提高到O(n log n) 甚至 O(n) 而不是当前较慢的 O(n^2) 解决方案(以及@Mazen 提出的解决方案)。 O(n^2) 在这里效率不高，因为它会导致大约 10,000,000 * 10,000,000 = 1000,000 亿次操作，这对于任何现代计算机来说都太多了。

这是纯 Python 中的哈希表解决方案:

table = {tuple(A[i]):i for i in range(A.shape[0])}
extract_BinA = np.empty(B.shape[0])
for i in range(B.shape[0]):
    val = tuple(B[i])
    if val in table:
        extract_BinA[i] = table[val]

请注意，如果 A 中的同一位置有多个点，结果可能会有所不同。

这是一个包含两个大小为 10,000 的随机数组的基准测试:

Initial solution: 53.82 s
Mazen solution:    1.76 s
This solution:     0.02 s

在这个小输入上，上述代码比初始解决方案快 2700 倍，比建议的替代解决方案快 88 倍。输入越大，差距就会越大，上面的代码比其他两个解决方案快很多数量级(即快 >10000 倍)。

---

更新:

如果 A 中有多个点彼此相等，则可以修改字典以存储索引列表而不是一个值。或者，可以创建字典，以便像在原始代码中一样保留第一个值。以下是两种解决方案的示例:

table = dict()
for i in range(A.shape[0])
    key = tuple(A[i])
    if key in table:
        table[key].append(i)
    else:
        table[key] = [i]

extract_BinA = np.empty(B.shape[0])
for i in range(B.shape[0]):
    val = tuple(B[i])
    if val in table:
        # Here table[val] is a list and thus you 
        # can do whatever you want with the indices. 
        # For example you can take the first one like here, 
        # or possibly the last as you want.
        extract_BinA[i] = table[val][0]

# Select always directly the first index
table = dict()
for i in range(A.shape[0])
    key = tuple(A[i])
    if key not in table:
        table[key] = i

extract_BinA = np.empty(B.shape[0])
for i in range(B.shape[0]):
    val = tuple(B[i])
    if val in table:
        extract_BinA[i] = table[val]

请注意，这些解决方案比上面的代码慢一点，但复杂度仍然是线性的(因此仍然非常快)。