python - 按行在另一个数组中搜索一个数组的元素 - Python/NumPy

Question

例如，我有一个唯一元素矩阵，

a=[
    [1,2,3,4],
    [7,5,8,6]
]

还有另一个独特的矩阵，其中填充了第一个矩阵中出现的元素。

b=[
    [4,1],
    [5,6]
]

我期待

的结果

[
    [3,0],
    [1,3]
].

也就是说，我想找到b的每一行元素与同一行中a的某些元素相等，返回这些元素在a中的索引。我怎样才能做到这一点？谢谢。

Answer 1

这是一个向量化的方法-

# https://stackoverflow.com/a/40588862/ @Divakar
def searchsorted2d(a,b):
    m,n = a.shape
    max_num = np.maximum(a.max() - a.min(), b.max() - b.min()) + 1
    r = max_num*np.arange(a.shape[0])[:,None]
    p = np.searchsorted( (a+r).ravel(), (b+r).ravel() ).reshape(m,-1)
    return p - n*(np.arange(m)[:,None])

def search_indices(a, b):
    sidx = a.argsort(1)
    a_s = np.take_along_axis(a,sidx,axis=1)
    return np.take_along_axis(sidx,searchsorted2d(a_s,b),axis=1)

sample 运行-

In [54]: a
Out[54]: 
array([[1, 2, 3, 4],
       [7, 5, 8, 6]])

In [55]: b
Out[55]: 
array([[4, 1],
       [5, 6]])

In [56]: search_indices(a, b)
Out[56]: 
array([[3, 0],
       [1, 3]])

另一个矢量化利用 broadcasting -

In [65]: (a[:,None,:]==b[:,:,None]).argmax(2)
Out[65]: 
array([[3, 0],
       [1, 3]])