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. 获取[0,0,0]出现在给定3D数组中的索引。
2. 将给定的3D数组从[0,0,0]首先出现的位置切片。

换句话说，我试图从给定的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]])

但是，如果这样做，我就会得到以下错误：

无法从形状广播输入数组..。变成形状..。

我期待着这样的事情：

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