python分割成数组_python numpy以0为边界分割数组

在做项目的时候需要分割数组，但是numpy本身并不提供根据某个数值对数组进行分割的函数，就类似string中的split能达到的功能。所以研究了一下。

例如：

我想做到把数组 [ 0, 1, 1, 0, 1, 3, 6, 3, 0, 1 ] 分割成 [ [1, 1], [1, 3, 6, 3], [1] ]

也就是以0为分界线分割数组，把非0数组单独分成一个数组，那么我可以这么做：

import numpy as np

a = np.array([0, 1, 1, 0, 1, 3, 6, 3, 0, 1])

pos = np.where(a > 0)[0]

split = np.where(np.diff(pos) != 1)[0] + 1

arr = np.split(a[pos], split)

如果需要修改成以其他数值为分界线，那么直接修改np.where()中的数值就行。

之后可以延伸出其他功能，

例如：

我想做到把数组 [ 0, 1, 1, 0, 1, 3, 6, 3, 0, 1 ] 分割成 [ [0], [1, 1], [0], [1, 3, 6, 3], [0], [1] ]

也就是说同样是以0位分割点进行分割，但是同时保留0，可以这么做：

import numpy as np

a = np.array([0, 0, 0, 1, 3, 4, 0, 1, 3, 4, 0])

b = np.where(np.diff(a != 0))[0] + 1

arr = np.split(a, b)