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

Question

我需要找到一种方法来执行此操作，我有一个形状为

[ batch_size , 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]]])

- RolandDeschain

---

你应该始终添加一些可用于测试的数据 - 以便我们更容易地进行测试。 - wwii

Answer 1

这里提供一种矢量化的方法 -

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

样例运行 -

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

- Divakar

---

嗨，谢谢快速回答。这给了我索引，对吧？针对每个元素，我如何以向量化的方式获取这些索引的值。我需要的是 a[0,b[0],0]，a[1,b[1],0]，等等... - RolandDeschain

@RolandDeschain 使用a[np.arange(a.shape[0])[:,None],idx,0]，其中idx是从发布的解决方案中获取的索引。 - Divakar

谢谢！索引让我疯了哈哈。你知道哪些页面或资源可以让我学习更多关于复杂的numpy索引吗？ - RolandDeschain

@RolandDeschain 这可能是一个停靠点 - https://docs.scipy.org/doc/numpy-1.10.1/reference/arrays.indexing.html#advanced-indexing - Divakar