Find last nonzero element 3D array - numpy array

Question

I need to find a way to perform this operation, i have an array of shape

[ batch_size , 150 , 1 ]

representing batch_size sequences of integers, each sequence is 150 elements long, but in each one there are a lot of zeros that were added in order to make all the sequences of the same length. I need to find the last nonzero element for each sequence, and put it in an array, the shape of this array has to be [ batch_size ] . I'd try the following :

last = []
for j in range(0 , inputs.shape[0] ):
  tnew = np.array( inputs[j][:][0] )
  tnew = np.trim_zeros(tnew )
   last.append( int(tnew[-1]) )

but I don't know if there is a better way to do this, with out having to loop over each element like that.

Thanks for your answers and help.

---

Test data

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

Answer 1

Here's one vectorized approach -

a.shape[1] - (a!=0)[:,::-1].argmax(1) - 1

Sample run -

In [191]: a = np.random.randint(0,3,(3,6,1))

In [192]: a
Out[192]: 
array([[[2],
        [1],
        [2],
        [2],
        [2],
        [0]],

       [[2],
        [1],
        [1],
        [0],
        [2],
        [0]],

       [[2],
        [1],
        [2],
        [0],
        [1],
        [1]]])

In [193]: a.shape[1] - (a!=0)[:,::-1].argmax(1) - 1
Out[193]: 
array([[4],
       [4],
       [5]])