python - 从 N 维 Numpy 数组中仅获取非零子数组

Question

我有一个 numpy array 'arr' 是 shape (1756020, 28, 28, 4)。基本上 'arr' 有 1756020 个 shape (28,28,4) 的小数组。在 1756020 数组中，967210 都是“全零”，而 788810 具有所有非零值。我想删除所有 967210“全零”小数组。我使用条件 arr[i]==0.any() 编写了一个 if else 循环，但这需要很多时间。有更好的方法吗？

Answer 1

向量化逻辑的一种方法是使用 numpy.any axis 的元组参数包含未测试的维度。

# set up 4d array of ones
A = np.ones((5, 3, 3, 4))

# make second of shape (3, 3, 4) = 0
A[1] = 0  # or A[1, ...] = 0; or A[1, :, :, :] = 0

# find out which are non-zero
res = np.any(A, axis=(1, 2, 3))

print(res)

[True False True True True]

此功能在 numpy v0.17 及更高版本中可用。根据 docs :

axis : None or int or tuple of ints, optional

If this is a tuple of ints, a reduction is performed on multiple axes, instead of a single axis or all the axes as before.