如何在numpy数组的每列中找到第一个非零值?

53

假设我有一个形如numpy数组:

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

我想找出每一列第一个非零值的索引。

因此,在这种情况下,我希望返回以下内容:

[0,0,2]

我该怎么做呢?

2个回答

107

首次出现的索引

在非零的掩码上沿着该轴(这里是列的第零轴)使用np.argmax,以获取首次匹配(True 值)的索引 -

(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,因此如果所需的invalid_val0,则我们可以直接使用mask.argmax(axis=axis)获取最终输出。
示例运行 -
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])

扩展以涵盖所有比较操作

要找到第一个zeros,只需将arr == 0用作函数中的mask。对于第一个等于某个值valones,使用arr == val,对于此处可能出现的所有比较情况都是如此。


最后出现的索引

要找到与某个比较条件相匹配的最后一个元素,我们需要沿着该轴翻转并使用与使用argmax相同的思路,然后通过从轴长度进行偏移来补偿翻转,如下所示 -

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])

再次强调,使用相应的比较器获取mask,然后在列出的函数中使用,即可涵盖所有可能的comparisons情况。


1
即使找到了所有第一个非零值,argmax 仍然会不必要地浏览数组的其余部分(可能很大),假设可能在那里找到更大的值(不知道 mask 没有比 1 更大的值)。是否可以轻松避免这种情况,而不需要“手动”实现板块处理? - root
1
@root 我有同样的问题。事实证明,自2012年以来,他们一直在讨论这个问题,并且还没有达成简单的解决方案! - Bill
应该是“某个比较标准”。 - KeithB

6
显然是二维问题,可以通过对每行应用找到第一个非零元素的函数(与问题中完全相同)来解决。
arr = np.array([[1,1,0],[1,1,0],[0,0,1],[0,0,0]])

def first_nonzero_index(array):
    """Return the index of the first non-zero element of array. If all elements are zero, return -1."""
    
    fnzi = -1 # first non-zero index
    indices = np.flatnonzero(array)
       
    if (len(indices) > 0):
        fnzi = indices[0]
        
    return fnzi

np.apply_along_axis(first_nonzero_index, axis=1, arr=arr)

# result
array([ 0,  0,  2, -1])

解释

np.flatnonzero(array) 方法(如 Henrik Koberg 在评论中所建议的)返回“数组扁平化后非零索引”。该函数计算这些索引并返回第一个(如果所有元素都为零则返回-1)。

apply_along_axis 方法沿指定轴应用函数到1-D切片。在这里,由于 axis 为 1,因此该函数应用于行。

如果我们可以假设输入数组的所有行至少都包含一个非零元素,则可以将解决方案写成一行代码:

np.apply_along_axis(lambda a: np.flatnonzero(a)[0], axis=1, arr=arr)

可能的变化

  • 如果我们对最后一个非零元素感兴趣,可以在函数中将indices [0]更改为indices [-1]。
  • 要按行获取第一个非零值,我们将在np.apply_along_axis中将axes = 1更改为axis = 0。

原始答案

这里是另一种使用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

1
不错的解决方案!argwhere这个名称相当不直观。 - aksg87
1
我不确定这个答案如何适用于涉及二维数组的问题。此外,在这里,您可以只使用 flatnonzero(array) 而不是 argwhere(array).squeeze() - Henrik Koberg
1
谢谢 Henrik, 我不知道 flatnonzero() 方法,我会用它来完善这个答案。关于问题的维度:数组确实是二维的,但所要求的问题是一维的:“找到每行中第一个非零元素”。这意味着可以将一维解决方案应用于每一行,并获得二维解决方案。 - MarcoP
1
同意,Marco。最后一点:虽然这个解决方案对于较小的数组效果很好,但是对于较大的数组来说性能会非常差。这是因为apply_along_axis是一个非向量化的便捷函数。在这种情况下,我宁愿选择被接受的答案。 - Henrik Koberg
我完全同意你的观点。 - MarcoP
这个代码运行了一个for循环,没有进行向量化处理。 - gsandhu

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接