如何找到第一个非零值在每一列的Numpy数组？

Question

假设我有一个如下形式的numpy数组：

arr=numpy.array([[1,1,0],[1,1,0],[0,0,1],[0,0,0]])

我想找到第一个索引的索引（对于每一列），其中的值是非零的。

因此，在本例中，我希望返回以下内容：

[0,0,2]

我该怎么办？

Answer 1

下面是一个使用numpy的替代方法。argwhere返回数组中非零元素的索引：

array = np.array([0,0,0,1,2,3,0,0])

nonzero_indx = np.argwhere(array).squeeze()
start, end = (nonzero_indx[0], nonzero_indx[-1])
print(array[start], array[end])

给予：

1 3

Answer 2

使用np。argmax沿非零掩码上的该轴（此处列的第零轴）获取第一个匹配的索引（真值）-

(arr!=0).argmax(axis=0)

扩展到覆盖通用轴说明符，对于在元素的轴上没有非零的情况，我们会有这样的实现-

def first_nonzero(arr, axis, invalid_val=-1):
    mask = arr!=0
    return np.where(mask.any(axis=axis), mask.argmax(axis=axis), invalid_val)

请注意，由于所有False值上的argmax（）返回0，因此如果所需的无效值0，我们将直接使用掩码获得最终输出。argmax（轴=轴）。

样本运行-

In [296]: arr    # Different from given sample for variety
Out[296]: 
array([[1, 0, 0],
       [1, 1, 0],
       [0, 1, 0],
       [0, 0, 0]])

In [297]: first_nonzero(arr, axis=0, invalid_val=-1)
Out[297]: array([ 0,  1, -1])

In [298]: first_nonzero(arr, axis=1, invalid_val=-1)
Out[298]: array([ 0,  0,  1, -1])

扩展到涵盖所有比较操作

要查找第一个零，只需在函数中使用arr==0作为mask。对于第一个值等于某个值的值val，在此处所有可能的比较情况下使用arr==val，依此类推。

为了找到匹配某个比较标准的最后一个，我们需要沿着该轴翻转，并使用使用artmax的相同想法，然后通过从轴长度偏移来补偿翻转，如下所示-

def last_nonzero(arr, axis, invalid_val=-1):
    mask = arr!=0
    val = arr.shape[axis] - np.flip(mask, axis=axis).argmax(axis=axis) - 1
    return np.where(mask.any(axis=axis), val, invalid_val)

样本运行-

In [320]: arr
Out[320]: 
array([[1, 0, 0],
       [1, 1, 0],
       [0, 1, 0],
       [0, 0, 0]])

In [321]: last_nonzero(arr, axis=0, invalid_val=-1)
Out[321]: array([ 1,  2, -1])

In [322]: last_nonzero(arr, axis=1, invalid_val=-1)
Out[322]: array([ 0,  1,  1, -1])

同样，这里所有可能的比较的情况都可以通过使用相应的比较器来获取掩码，然后在列出的函数中使用。