（Python）找出两个数组中与另外两个数组中的值相等的索引

Question

我可以帮助你翻译以下关于编程的内容：我有以下4个数组，我想获取与B和Y中相同位置的值对应的在A和X中相等的值的索引。因此，对于以下示例：

    import numpy as np
    A = np.asarray([400.5, 100,  700,   200,  15, 900])
    B = np.asarray([500.5, 200,  500, 600.5,   8, 999])
    X = np.asarray([400.5, 700,  100,   300,  15, 555, 900])
    Y = np.asarray([500.5, 500,600.5,   100,   8, 555, 999])

我想获得两个包含索引的数组：

indAB = [0 2 4 5]

• 0，因为A&B中的400.5和500.5在X&Y中的位置0也存在
• 2，因为A&B中的700和500在X&Y中的位置2也存在
• 4，因为A&B中的15和8在X&Y中的位置4也存在
• 5，因为A&B中的900和999在X&Y中的位置5也存在

indXY = [0 1 4 6]

• 0、1、4和6与indAB相似，但是是相对于X&Y而言的。

其中indAB是A和B中与X和Y相等的值的索引，indXY是X和Y中与A和B相等的值的索引。

这是我目前的尝试：

    def indices(a,b):
        setb = set(b)
        ind = [i for i, x in enumerate(a) if x in setb]
        return ind

    iA = np.asarray(indices(A,X))
    iB = np.asarray(indices(X,A))
    iX = np.asarray(indices(B,Y))
    iY = np.asarray(indices(Y,B))

    def CommonIndices(a,b):
        return np.asarray(list(set(a) & set(b)))

    indAB = CommonIndices(iA,iX)
    indXY = CommonIndices(iB,iY)

    print(indAB) # returns = [0 2 4 5]
    print(indXY) # returns = [0 1 2 4 6]

我一直得到[0 1 2 4 6]这个错误的indXY。2不应该被包括在内，因为即使600.5在Y和B中，A和B中的200和100（分别）也不相等。

如果有人能提供解决方案，我将非常感激。非常感谢！

- TimeExplorer

Answer 1

numpy_indexed 包（声明: 我是它的作者）包含了高效且优雅的方法来实现这种操作。内存要求是线性的，计算需求为 NlogN。对于您考虑的大型数组，与目前被接受的暴力方法相比，速度提升可以轻易地达到数量级。

import numpy as np
import numpy_indexed as npi

A = np.asarray([400.5, 100,  700,   200,  15, 900])
B = np.asarray([500.5, 200,  500, 600.5,   8, 999])
X = np.asarray([400.5, 700,  100,   300,  15, 555, 900])
Y = np.asarray([500.5, 500,600.5,   100,   8, 555, 999])

AB = np.stack([A, B], axis=-1)
XY = np.stack([X, Y], axis=-1)

# casting the AB and XY arrays to npi.index first is not required, but a performance optimization; without this each call to npi.indices would have to re-index the arrays, which is the expensive part
AB = npi.as_index(AB)
XY = npi.as_index(XY)
# npi.indices(list, items) is a vectorized nd-equivalent of list.index(item)
indAB = npi.indices(AB, XY, missing='mask').compressed()
indXY = npi.indices(XY, AB, missing='mask').compressed()

请注意，您还可以选择如何处理缺失值。还要看一下集合操作，比如npi.intersection(XY, AB); 在更高的层次上，它们可能提供了更简单的方法来实现您想要达到的目标。

- Eelco Hoogendoorn

---

不错。它确实比被接受的答案更有效率。我不知道你是如何实现 npi 的。据我所知，我的答案与你的相关，但使用了原始的Python对象。 - Eric Duminil

有点不同；npi是用“纯numpy”编写的，因此执行集合类型操作的技巧是对数组进行（arg）排序并将相关项放在一起。因此，性能为O（NlogN），而不是您的方法应该具有的O（N）。但是对于许多实际情况，排序比O（NlogN）更好，因为数据很少是纯随机的；当然，向量化很难被击败。 - Eelco Hoogendoorn

numpy_indexed包运行得非常完美。非常感谢！:D - TimeExplorer

@EelcoHoogendoorn：非常感谢您的回答。与SO上的作者进行第一手经验交流总是很愉快的。 - Eric Duminil

Answer 2

试试这个：

import numpy as np

A = np.asarray([400.5, 100,  700,   200,  15, 900])
B = np.asarray([500.5, 200,  500, 600.5,   8, 999])
X = np.asarray([400.5, 700,  100,   300,  15, 555, 900])
Y = np.asarray([500.5, 500,600.5,   100,   8, 555, 999])

AB = np.stack([A, B], axis=-1)
XY = np.stack([X, Y], axis=-1)

eq = AB[:, np.newaxis, :] == XY[np.newaxis, :, :]
eq = np.logical_and.reduce(eq, axis=-1)

indAB, = np.where(np.logical_or.reduce(eq, axis=1))
indXY, = np.where(np.logical_or.reduce(eq, axis=0))

print("indAB", indAB)
print("indXY", indXY)

输出：

indAB [0 2 4 5]
indXY [0 1 4 6]

Explanation

AB和XY只是将A和B以及X和Y这两个数组“堆叠”成二维数组。 eq保存了AB和XY中元素的全组合比较结果；使用np.newaxis为AB和XY添加了维度（注意，AB在第1个位置上得到了一个新的维度，而XY在位置0上得到了一个新的维度）。等号运算符==通过它们的新维度广播数组。第一个np.logical_and.reduce是确保“两个”“元素”都相等（A与X且B与Y），np.logical_or.reduce操作检查是否有任何来自AB到XY和从XY到AB的完全相等性。最后，np.where获取了索引。

作为缺点，请注意这需要一个大小为len（A） x len（X） x 2的布尔数组，所以如果原始数组非常大，您可能会遇到内存问题。 更新 如上所述，非常大的数组可能是一个问题。 如果您想要“一次性”进行所有比较，则实际上没有绕过它的方法（中间数组的大小就是比较的数量）。 但是，您也可以“分段”运行算法，例如像这样：

import numpy as np

MAX_SIZE = 2  # Biggest array will be MAX_SIZE x MAX_SIZE x 2

A = np.asarray([400.5, 100,  700,   200,  15, 900])
B = np.asarray([500.5, 200,  500, 600.5,   8, 999])
X = np.asarray([400.5, 700,  100,   300,  15, 555, 900])
Y = np.asarray([500.5, 500,600.5,   100,   8, 555, 999])

AB = np.stack([A, B], axis=-1)
XY = np.stack([X, Y], axis=-1)

maskAB = np.full(len(AB), False, dtype=bool)
maskXY = np.full(len(XY), False, dtype=bool)

for iAB in range(0, len(AB), MAX_SIZE):
    pAB = np.expand_dims(AB[iAB:iAB + MAX_SIZE], axis=1)
    for iXY in range(0, len(XY), MAX_SIZE):
        pXY = np.expand_dims(XY[iXY:iXY + MAX_SIZE], axis=0)
        eq = pAB == pXY
        eq = np.logical_and.reduce(eq, axis=-1)
        maskAB[iAB:iAB + MAX_SIZE] |= np.logical_or.reduce(eq, axis=1)
        maskXY[iXY:iXY + MAX_SIZE] |= np.logical_or.reduce(eq, axis=0)

indAB, = np.where(maskAB)
indXY, = np.where(maskXY)

print("indAB", indAB)
print("indXY", indXY)

而输出仍然是：

indAB [0 2 4 5]
indXY [0 1 4 6]

我只是使用了MAX_SIZE为2，仅仅是为了说明在示例中它能够工作，但在实际应用中，您可以根据您愿意使用的最大内存量来选择它（例如，对于MAX_SIZE = 10000，它应该在数百兆字节的数量级上）。 MAX_SIZE不需要小于数组的大小，也不必是其大小的除数。

- jdehesa

---

非常感谢！！这正是我所需要的，我自己永远也达不到这个水平哈哈。非常感谢:D - TimeExplorer

@TimeExplorer 没问题。我已经添加了一些解释，以防您（或任何找到答案的人）发现它有用。 - jdehesa

这些数组确实相当大，3518220442072=143838581348。有没有一种不太昂贵的方法可以解决这个问题？ - TimeExplorer

@TimeExplorer，我添加了一种运行算法“分段”的代码变体，它应该允许您限制所使用的内存量。 - jdehesa

Answer 3

这里有一种替代方法。我敢说它比较清晰，由于使用了集合，应该很有效，并且只需要O(len(A)+len(X))的内存。 numpy甚至不需要，但可以用于数组。

from collections import defaultdict

A = [400.5, 100, 700, 200, 15, 900]
B = [500.5, 200, 500, 600.5, 8, 999]
X = [400.5, 700, 100, 300, 15, 555, 900]
Y = [500.5, 500, 600.5, 100, 8, 555, 999]

def get_indices(values):
    d = defaultdict(set)
    for i, value in enumerate(values):
        d[value].add(i)
    return d

iA, iB, iX, iY = [get_indices(values) for values in [A, B, X, Y]]
print(iA)
# {400.5: {0}, 100: {1}, 200: {3}, 900: {5}, 700: {2}, 15: {4}}
print(iX)
# {400.5: {0}, 100: {2}, 300: {3}, 900: {6}, 555: {5}, 700: {1}, 15: {4}}

for i, (a, b) in enumerate(zip(A, B)):
    common_indices = iX[a] & iY[b]
    if common_indices:
        print("A B : %d" % i)
        print("X Y : %d" % common_indices.pop())
        print()

#   A B : 0
#   X Y : 0

#   A B : 2
#   X Y : 1

#   A B : 4
#   X Y : 4

#   A B : 5
#   X Y : 6

- Eric Duminil