（Python）查找两个数组中的值的索引，这些值与另两个数组中的值相等

Question

我有以下4个数组，我想得到数组A和X中相等的值的索引，这些值对应于B和Y中相同位置的值。所以对于下面的例子

    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也在位置2的X&Y中
• 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中的值

这是我目前为止的尝试：

^{pr2}$

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

如果有人能提出解决办法，我将不胜感激。非常感谢！在

Answer 1

这是另一种方法。我敢说这是相对清楚的，由于使用了集合，它应该是高效的，而且它只需要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

Answer 2

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()

请注意，您也可以选择如何处理缺少的值。还要看一下set操作，比如npi交叉口（XY，AB）；它们可能提供一个更简单的途径，让你在更高的层次上实现你的目标。在

Answer 3

试试这个：

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)

输出：

^{pr2}$

说明

AB和{}只是数组A和{}和{}和{}分别“堆叠”成二维数组。eq保存AB和XY中元素的all-overall比较；np.newaxis用于向AB和{}添加维度（注意，AB在位置1获得一个新维度，在位置0中{}获得一个新维度）。相等运算符==通过数组的新维度广播数组。第一个np.logical_and.reduce是确保这两个的“组件”是相等的（A到{}和{}到{}），并且np.logical_or.reduce操作检查从AB到{}和从XY到{}之间是否存在完全相等。最后，np.where得到索引。在

作为缺点，请注意，这需要一个大小为len(A)xlen(X)x2的布尔数组，因此，如果原始数组非常则可能会遇到内存问题。在

更新

如前所述，非常大的阵列可能是一个问题。如果您想“一次性”进行所有的比较，实际上并没有办法绕过它（中间数组的大小就是比较的次数）。但是，您也可以“按块”运行算法，例如：

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)

输出仍然是：

^{pr2}$

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