获取NumPy二维数组交叉行的索引

Question

我想要找到一个主要的二维数组 A 和另一个数组 B 中相交的行的索引。

A=array([[1, 2],
         [3, 4],
         [5, 6],
         [7, 8],
         [9, 10]])

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

result=[0,2]

这个操作应该返回 [0,2]，这是基于数组 A 的索引。

那么，如何高效地在二维数组中做到这一点呢？

谢谢！

补充说明

我尝试过这个函数：

k[np.in1d(k.view(dtype='i,i').reshape(k.shape[0]),k2.view(dtype='i,i').
reshape(k2.shape[0]))]

来自 实现 numpy in1d 用于二维数组？，但是我遇到了重塑错误。我的数据类型是浮点数（保留两位小数）。此外，我也尝试过使用集合，但性能相当慢。

Answer 1

你可以使用 np.char.array() 这个东西来进行比较，方法是用 np.in1d()：

s1 = np.char.array(A[:,0]) + '-' + np.char.array(A[:,1])
s2 = np.char.array(B[:,0]) + '-' + np.char.array(B[:,1])

np.where(np.in1d(s1, s2))[0]
#array([0, 2], dtype=int64)

注意：A 和 B 必须是相同的数据类型（比如 int，float 等等），这样才能正常工作。

Answer 2

只需做一些小改动，你就可以让你的方法正常工作：

In [15]: A
Out[15]: 
array([[ 1,  2],
       [ 3,  4],
       [ 5,  6],
       [ 7,  8],
       [ 9, 10]])

In [16]: B
Out[16]: 
array([[1, 4],
       [1, 2],
       [5, 6],
       [6, 3]])

In [17]: np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1))
Out[17]: array([ True, False,  True, False, False], dtype=bool)

In [18]: np.nonzero(np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1)))
Out[18]: (array([0, 2], dtype=int64),)

In [19]: np.nonzero(np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1)))[0]
Out[19]: array([0, 2], dtype=int64)

如果你的数组不是浮点数，并且都是连续的，那么下面这个方法会更快：

In [21]: dt = np.dtype((np.void, A.dtype.itemsize * A.shape[1]))

In [22]: np.nonzero(np.in1d(A.view(dt).reshape(-1), B.view(dt).reshape(-1)))[0]
Out[22]: array([0, 2], dtype=int64)

这里还有一个快速的时间测试：

In [24]: %timeit np.nonzero(np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1)))[0]
10000 loops, best of 3: 75 µs per loop

In [25]: %timeit np.nonzero(np.in1d(A.view(dt).reshape(-1), B.view(dt).reshape(-1)))[0]
10000 loops, best of 3: 29.8 µs per loop