python - 使用 Numpy 索引 3D 数组并保持非零轴的优雅方法

Question

假设我有以下 3D 数组:

L=np.arange(18).reshape((2,3,3))
L[:,:,1] = 0; L[:,[0,1],:] = 0

In []: L
Out[]: 
array([[[ 0,  0,  0],
        [ 0,  0,  0],
        [ 6,  0,  8]],

       [[ 0,  0,  0],
        [ 0,  0,  0],
        [15,  0,  17]]])

L[0,:] 中的零列总是与 L[1,:] 中相应的零列匹配。

我现在想删除中间的列，其中沿轴的总和等于 0(忽略 行 为零。我目前笨拙的方法是

l=np.nonzero(L.sum(axis=1))[1]

In []: L[:,:,l[:len(l)/2]]
Out[]: 
array([[[ 0,  0],
        [ 0,  0],
        [ 6,  8]],

       [[ 0,  0],
        [ 0,  0],
        [15, 17]]])

执行此操作的迂回方式是什么？

Answer 1

我们可以沿着前两个轴寻找所有的零，并用它来屏蔽掉第三个轴上的那些 -

L[:,:,~(L==0).all(axis=(0,1))]

或者，使用 any() 替换 ~all() -

L[:,:,(L!=0).any(axis=(0,1))]

我们可以使用省略号 ... 来替换 :,: 并跳过 arg axis 来给我们一个紧凑的版本-

L[...,~(L==0).all((0,1))]
L[...,(L!=0).any((0,1))]

有关省略号如何用于 NumPy 数组的更多信息，here .

---

对于问题的总和部分，它会是类似的-

L[...,L.sum((0,1))!=0]