python - numpy数组中的向量化非零组

Question

假设您有 1d numpy 数组：

[0,0,0,0,0,1,2,3,0,0,0,0,4,5,0,0,0]

如果不使用 for 循环，您将如何创建以下组？

[1,2,3], [4,5]

标签： pythonnumpy

Answer 1

这是一种使用方法np.split：

a
# array([0, 0, 0, 0, 0, 1, 2, 3, 0, 0, 0, 0, 4, 5, 0, 0, 0])
### find nonzeros
z = a!=0
### find switching points
z[1:] ^= z[:-1]
### split at switching points and discard zeros
np.split(a, *np.where(z))[1::2]
# [array([1, 2, 3]), array([4, 5])]