python - Numpy 搜索和切片 3D 数组

Question

我是 Python 和 Numpy 的新手，正在尝试完成以下任务:

给定的 3D 数组:

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

1. 获取给定 3D 数组中 [0,0,0] 出现的索引。
2. 从 [0,0,0] 首先出现的位置对给定的 3D 数组进行切片。

换句话说，我试图从给定的 3D 数组中删除填充(在本例中为:[0,0,0])。

这是我尝试过的，

arr_zero = np.zeros(3)
for index in range(0, len(arr_3d)):
    rows, cols = np.where(arr_3d[index] == arr_zero)
    arr_3d[index] = np.array(arr_3d[0][:rows[0]])

但是这样做，我不断收到以下错误:

Could not broadcast input array from shape ... into shape ...

我期待这样的事情:

[[[1,2,3],[4,5,6]],
 [[3,2,1]]
 [[1,2,3],[4,5,6],[7,8,9]]]

如有任何帮助，我们将不胜感激。

Answer 1

使用 all() 缩减和 argmax() 获取这些索引的第一次出现，然后将每个 2D 切片从 3D 数组 -

In [106]: idx = (arr_3d == [0,0,0]).all(-1).argmax(-1)

# Output as list of arrays
In [107]: [a[:i] for a,i in zip(arr_3d,idx)]
Out[107]: 
[array([[1, 2, 3],
        [4, 5, 6]]), array([[3, 2, 1]]), array([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])]

# Output as list of lists
In [108]: [a[:i].tolist() for a,i in zip(arr_3d,idx)]
Out[108]: [[[1, 2, 3], [4, 5, 6]], [[3, 2, 1]], [[1, 2, 3], [4, 5, 6], [7, 8, 9]]]