基于一列的交集过滤多个NumPy数组

Question

我有三个相当大的NumPy数组，具有不同的行数，它们的第一列都是integers。我希望对这些数组进行筛选，以便只剩下第一列中的值由这三个数组共享的行。这将留下三个大小相同的数组。其他列中的条目不一定在数组之间共享。在

因此，输入：

A = 
[[1, 1],
[2, 2],
[3, 3],]

B = 
[[2, 1],
[3, 2],
[4, 3],
[5, 4]]

C = 
[[2, 2],
[3, 1]
[5, 2]]

我希望作为输出返回：

^{pr2}$

我目前的做法是：

1. 使用numpy.intersect1d()

2. 在此交集和每个数组的第一列上使用numpy.in1d()来查找每个数组中未共享的行索引（使用此处找到的方法的修改版本将boolean转换为索引：Python: intersection indices numpy array）

3. 最后将numpy.delete()与这些索引及其各自的数组一起使用，以删除第一列中包含非共享项的行。

我想知道是否有一种更快或更优雅的Python式的方法来实现这一点，但它适合于非常大的数组。在

Answer 1

一种方法是构建一个指示符数组，或者根据需要创建一个哈希表，以指示所有输入数组中的整数。然后您可以使用基于此指示器数组的布尔索引来获得子数组。像这样：

import numpy as np

# Setup
A = np.array(
[[1, 1],
[2, 2],
[3, 3],])

B = np.array(
[[2, 1],
[3, 2],
[4, 3],
[5, 4]])

C = np.array(
[[2, 2],
[3, 1],
[5, 2],])


def take_overlap(*input):
    n = len(input)
    maxIndex = max(array[:, 0].max() for array in input)
    indicator = np.zeros(maxIndex + 1, dtype=int)
    for array in input:
        indicator[array[:, 0]] += 1
    indicator = indicator == n

    result = []
    for array in input:
        # Look up each integer in the indicator array
        mask = indicator[array[:, 0]]
        # Use boolean indexing to get the sub array
        result.append(array[mask])

    return result

subA, subB, subC = take_overlap(A, B, C)

这应该很快，而且这个方法不假设输入数组的元素是唯一的或已排序的。然而，如果索引整数是稀疏的，即[1，10，10000]，这种方法可能会占用大量内存，并且速度可能会稍慢一些，但是如果这些整数的密度或多或少，则应该接近最优。在

Answer 2

这是可行的，但我不确定它是否比其他任何答案都快：

import numpy as np

A = np.array(
[[1, 1],
[2, 2],
[3, 3],])

B = np.array(
[[2, 1],
[3, 2],
[4, 3],
[5, 4]])

C = np.array(
[[2, 2],
[3, 1],
[5, 2],])

a = A[:,0]
b = B[:,0]
c = C[:,0]

ab = np.where(a[:, np.newaxis] == b[np.newaxis, :])
bc = np.where(b[:, np.newaxis] == c[np.newaxis, :])

ab_in_bc = np.in1d(ab[1], bc[0])
bc_in_ab = np.in1d(bc[0], ab[1])

arows = ab[0][ab_in_bc]
brows = ab[1][ab_in_bc]
crows = bc[1][bc_in_ab]

anew = A[arows, :]
bnew = B[brows, :]
cnew = C[crows, :]

print(anew)
print(bnew)
print(cnew)

给出：

^{pr2}$

Answer 3

示例中的索引是经过排序且唯一的。假设这不是巧合（这种情况经常发生，或很容易实施），则以下工作：

import numpy as np

A = np.array(
[[1, 1],
[2, 2],
[3, 3],])

B = np.array(
[[2, 1],
[3, 2],
[4, 3],
[5, 4]])

C = np.array(
[[2, 2],
[3, 1],
[5, 2],])

I = reduce(
    lambda l,r: np.intersect1d(l,r,True),
    (i[:,0] for i in (A,B,C)))

print A[np.searchsorted(A[:,0], I)]
print B[np.searchsorted(B[:,0], I)]
print C[np.searchsorted(C[:,0], I)]

如果第一列不是按排序顺序排列的（但仍然是唯一的）：

^{pr2}$

确保在

I = reduce(
    lambda l,r: np.intersect1d(l,r,False),
    (i[:,0] for i in (A,B,C)))

可以使用np.唯一在