在2d數組行中找到多個1d數組元素

Question

我想在2d數組的行中找到1d數組的元素。

例

In [1]: import numpy as np

In [2]: a = np.array([7,7,7])

In [3]: a
Out[3]: array([7, 7, 7])

In [4]: b = np.arange(15).reshape(3,5)

In [5]: b
Out[5]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

行0和2 b不具備的相應元素相匹配a ，但行1具有在位置2的預期輸出的匹配：

array([nan, 2, nan])

如果給定行中有多個匹配項，則應使用第一個匹配項的位置。

我可以使用普通的python循環研討解決方案，但我對矢量化的方法感興趣。

Answer 1

將b和一個元素進行明智的比較，然后為每一行找到第一個True值索引，如果所有False，則將其設置為nan。

np.where(np.sum(b==a[:,None],1) > 0, np.argmax(b==a[:,None],1), np.nan)
Out[22]: array([ nan,   2.,  nan])