numpy:在2D数组的每一行中找到第一个值索引

时间:2014-05-07 15:49:46

标签: numpy

如何使用矢量化numpy函数找到2D数组每行中第一个值的索引?

例如,给定

I = numpy.array([1,1,1]
M = numpy.array([[1,2,3],[2,3,1],[3,1,2]])

输出应为:

array([0, 2, 1])

我可以用这样的列表理解来做到这一点:

[ numpy.where(M[i] == I[i])[0][0] for i in range(0, len(I)) ]

numpy等价物会是什么?

2 个答案:

答案 0 :(得分:2)

利用矢量化的可能性如下

coords = ((I[:, np.newaxis] == M) * np.arange(M.shape[1], 0, -1)[np.newaxis, :]).argmax(1)
any = (I[:, np.newaxis] == M).any(1)
coords = coords[any]

它通过将减少的计数器乘以每一行来消除同一行中几个感兴趣的值之间的歧义,使得第一次出现具有最高值。如果给定的行不包含指示的值,则将其从coords中删除。剩余的行(其中找到了相应的值)由any

索引

答案 1 :(得分:1)

我认为这些可能会一步一步地实现:

In [52]:

I = np.array([1,1,1])
#M = np.array([[1,2,3],[2,3,1],[3,1,2]])
M = np.array([[4,2,3],[2,3,4],[3,4,2]])
In [53]:

I1=I.reshape((-1,1))
In [54]:

M1=np.hstack((M, I1))
In [55]:

np.apply_along_axis(np.argmax, 1, (M1-I1)==0)
Out[55]:
array([3, 3, 3])

如果在M中找不到该号码,则生成的索引为M.shape[1]。由于结果是array int,因此在这些单元格中放置nan不是一种选择。但如果结果为-1,我们可能会考虑将result用于这些情况:

result[result==(M.shape[1])]=-1