根据一列的交集过滤多个 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]]

我希望作为输出返回：

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


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

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

我目前的做法是：

1. 使用numpy.intersect1d()查找前三列的交集

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

3. 最后将numpy.delete() 与每个索引及其各自的数组一起使用，以删除第一列中具有非共享条目的行。

我想知道是否有更快或更优雅的 Pythonic 方式来解决这个问题，但它适用于非常大的数组。

Answer 1

您的示例中的索引已排序且唯一。假设这不是巧合（并且这种情况经常出现，或者很容易实施），以下工作：

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

如果第一列没有排序（但仍然是唯一的）：

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

def index_by_first_column_entry(M, keys):
    colkeys = M[:,0]
    sorter = np.argsort(colkeys)
    index = np.searchsorted(colkeys, keys, sorter = sorter)
    return M[sorter[index]]

print index_by_first_column_entry(C, I)

并确保在

中将true更改为false

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

可以使用 np.unique 对重复值进行泛化

Answer 2

执行此操作的一种方法是构建一个指标数组，或者如果您愿意，可以构建一个哈希表，以指示您的所有输入数组中有哪些整数。然后您可以使用基于此指标数组的布尔索引来获取子数组。像这样的：

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 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],])

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)

给予：

[[2 2]
 [3 3]]
[[2 1]
 [3 2]]
[[2 2]
 [3 1]]