Set last non-zero element of each row to zero - NumPy

Question

I have an array A:

A = array([[1, 2, 3,4], [5,6,7,0] , [8,9,0,0]])

I want to change the last non-zero of each row to 0

A = array([[1, 2, 3,0], [5,6,0,0] , [8,0,0,0]])

how to write the code for any n*m numpy array? Thanks, S ;-)

Answer 1

Approach #1

One approach based on cumsum and argmax -

A[np.arange(A.shape[0]),(A!=0).cumsum(1).argmax(1)] = 0

Sample run -

In [59]: A
Out[59]: 
array([[2, 0, 3, 4],
       [5, 6, 7, 0],
       [8, 9, 0, 0]])

In [60]: A[np.arange(A.shape[0]),(A!=0).cumsum(1).argmax(1)] = 0

In [61]: A
Out[61]: 
array([[2, 0, 3, 0],
       [5, 6, 0, 0],
       [8, 0, 0, 0]])

Approach #2

One more based on argmax and hopefully more efficient -

A[np.arange(A.shape[0]),A.shape[1] - 1 - (A[:,::-1]!=0).argmax(1)] = 0

---

Explanation

One of the uses of argmax is to get ID of the first occurence of the max element along an axis in an array . In the first approach we get the cumsum along the rows and get the first max ID, which represents the last non-zero elem. This is because cumsum on the leftover elements won't increase the sum value after that last non-zero element.

Let's re-run that case in a bit more detailed manner -

In [105]: A
Out[105]: 
array([[2, 0, 3, 4],
       [5, 6, 7, 0],
       [8, 9, 0, 0]])

In [106]: (A!=0)
Out[106]: 
array([[ True, False,  True,  True],
       [ True,  True,  True, False],
       [ True,  True, False, False]], dtype=bool)

In [107]: (A!=0).cumsum(1)
Out[107]: 
array([[1, 1, 2, 3],
       [1, 2, 3, 3],
       [1, 2, 2, 2]])

In [108]: (A!=0).cumsum(1).argmax(1)
Out[108]: array([3, 2, 1])

Finally, we use fancy-indexing to use those as the column indices to set appropriate elements in A.

In the second approach, when we use argmax on the boolean array, we simply got the first occurence of True, which we used on a row-flipped version of the input array. As such, we would have the last non-zero elem in the original order. Rest of the idea there, is the same.