Splitting Numpy array based on value

后端 未结 6 1930
我寻月下人不归
我寻月下人不归 2020-12-17 19:05

Suppose I have this NumPy array:

a = np.array([0, 3, 5, 5, 0, 10, 14, 15, 56, 0, 12, 23, 45, 23, 12, 45, 
              0, 1, 0, 2, 3, 4, 0, 0 ,0])
<         


        
6条回答
  •  庸人自扰
    2020-12-17 19:48

    Here's a vectorized approach using np.where and np.split -

    idx = np.where(a!=0)[0]
    aout = np.split(a[idx],np.where(np.diff(idx)!=1)[0]+1)
    

    Sample run -

    In [23]: a
    Out[23]: 
    array([ 0,  3,  5,  5,  0, 10, 14, 15, 56,  0,  0,  0, 12, 23, 45, 23, 12,
           45,  0,  1,  0,  2,  3,  4,  0,  0,  0])
    
    In [24]: idx = np.where(a!=0)[0]
    
    In [25]: np.split(a[idx],np.where(np.diff(idx)!=1)[0]+1)
    Out[25]: 
    [array([3, 5, 5]),
     array([10, 14, 15, 56]),
     array([12, 23, 45, 23, 12, 45]),
     array([1]),
     array([2, 3, 4])]
    

提交回复
热议问题