How to find first non-zero value in every column of a numpy array?

Question

Suppose I have a numpy array of the form:

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

I want to find the indices of the first index (for every column) where the value is non-zero.

So in this instance, I would like the following to be returned:

[0,0,2]

How do I go about this?

Answer 1

Indices of first occurrences

Use np.argmax along that axis (zeroth axis for columns here) on the mask of non-zeros to get the indices of first matches (True values) -

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

Extending to cover generic axis specifier and for cases where no non-zeros are found along that axis for an element, we would have an implementation like so -

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

Note that since argmax() on all False values returns 0, so if the invalid_val needed is 0, we would have the final output directly with mask.argmax(axis=axis).

Sample runs -

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

Extending to cover all comparison operations

To find the first zeros, simply use arr==0 as mask for use in the function. For first ones equal to a certain value val, use arr == val and so on for all cases of comparisons possible here.

---

Indices of last occurrences

To find the last ones matching a certain comparison criteria, we need to flip along that axis and use the same idea of using argmax and then compensate for the flipping by offsetting from the axis length, as shown below -

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)

Sample runs -

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

Again, all cases of comparisons possible here are covered by using the corresponding comparator to get mask and then using within the listed function.

Answer 2

The problem, apparently 2D, can be solved by applying to the each row a function that finds the first non-zero element (exactly as in the question).

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

Explanation

The np.flatnonzero(array) method (as suggested in the comments by Henrik Koberg) returns "indices that are non-zero in the flattened version of array". The function calculates these indices and returns the first (or -1 if all elements are zero).

The apply_along_axis applys a function to 1-D slices along the given axis. Here since the axis is 1, the function is applied to the rows.

If we can assume that all rows of the input array contain at leas one non-zero element, the solution can be written calculated in one line:

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

Possible variations

• If we were interested in the last non-zero element, that could be obrained by changing indices[0] into indices[-1] in the function.
• To get the first non-zero by row, we would change axes=1 into axis=0 in np.apply_along_axis

ORIGINAL ANSWER

Here is an alternative using numpy.argwhere which returns the index of the non zero elements of an array:

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

gives:

1 3