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. 从[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]])

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

无法将输入数组从形状 ... 广播到形状 ...

我期待这样的事情：

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

任何帮助将不胜感激。

Answer 1

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

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