python - 查找最后一个非零元素 3D 数组 - numpy 数组

Question

我需要找到一种方法来执行这个操作，我有一个形状数组

[批量大小, 150, 1]

代表 batch_size 整数序列，每个序列有 150 个元素长，但在每个序列中都有很多添加的零，以使所有序列具有相同的长度。我需要为每个序列找到最后一个非零元素，并将其放入一个数组中，这个数组的形状必须是 [ batch_size ] 。我会尝试以下操作:

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

但我不知道是否有更好的方法来做到这一点，而不必像那样遍历每个元素。

感谢您的回答和帮助。

---

测试数据

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

这是一种矢量化方法 -

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

sample 运行-

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