python - 根据一列的交集过滤多个 NumPy 数组

Question

我有三个相当大的 NumPy 数组，它们的行数各不相同，第一列都是 整数。我的希望是过滤这些数组，以便只剩下第一列中的值由所有三个共享的行。这将留下三个相同大小的数组。其他列中的条目不一定跨数组共享。

因此，输入:

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 使用此处找到的方法的修改版本进行索引:Python: intersection indices numpy array )

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 泛化到重复值