How to split a numpy array based on a column?

后端 未结 3 1557
小鲜肉
小鲜肉 2021-01-24 18:45

I have an array of the form :

[[ 1. ,    2.,     3.,     1.,     3.,     3.,     4.   ],
 [ 1.3,    2.3,    3.3,    3.,     3.3,    3.3,    4.3  ],
 [ 1.2,    2.         


        
3条回答
  •  暗喜
    暗喜 (楼主)
    2021-01-24 19:34

    You can do this in O(NlogN) time using numpy.argsort, numpy.array_split, numpy.diff and numpy.where:

    >>> indices = np.argsort(arr[:, 3])
    >>> arr_temp = arr[indices]
    >>> np.array_split(arr_temp, np.where(np.diff(arr_temp[:,3])!=0)[0]+1)
    [array([[ 1.  ,  2.  ,  3.  ,  1.  ,  3.  ,  3.  ,  4.  ],
           [ 1.89,  2.3 ,  1.  ,  1.  ,  3.  ,  3.  ,  4.  ],
           [ 1.1 ,  2.1 ,  1.  ,  1.  ,  3.  ,  3.  ,  4.  ],
           [ 1.9 ,  2.2 ,  1.  ,  1.  ,  3.  ,  3.  ,  4.  ],
           [ 1.3 ,  2.2 ,  1.  ,  1.  ,  3.  ,  3.  ,  4.  ],
           [ 1.5 ,  2.1 ,  1.  ,  1.  ,  3.  ,  3.  ,  4.  ],
           [ 1.4 ,  2.3 ,  1.  ,  1.  ,  3.  ,  3.  ,  4.  ]]), array([[ 1.2  ,  2.8  ,  3.2  ,  2.   ,  3.66 ,  3.2  ,  4.2  ],
           [ 1.2  ,  2.7  ,  3.2  ,  2.   ,  3.2  ,  3.231,  4.2  ],
           [ 1.2  ,  2.9  ,  3.2  ,  2.   ,  3.2  ,  3.2  ,  4.2  ],
           [ 1.2  ,  2.9  ,  3.2  ,  2.   ,  3.34 ,  3.2  ,  4.2  ],
           [ 1.2  ,  2.8  ,  3.2  ,  2.   ,  3.2  ,  3.2  ,  4.2  ],
           [ 1.2  ,  2.7  ,  3.2  ,  2.   ,  3.2  ,  3.2  ,  4.2  ],
           [ 1.2  ,  2.2  ,  3.2  ,  2.   ,  3.2  ,  3.2  ,  4.2  ]]), array([[ 1.3 ,  2.3 ,  3.6 ,  3.  ,  3.3 ,  3.3 ,  4.3 ],
           [ 1.89,  2.3 ,  3.5 ,  3.  ,  3.3 ,  3.3 ,  4.3 ],
           [ 1.3 ,  2.3 ,  3.5 ,  3.  ,  3.3 ,  3.3 ,  4.3 ],
           [ 1.3 ,  2.22,  3.6 ,  3.  ,  3.3 ,  3.3 ,  4.3 ],
           [ 1.3 ,  2.3 ,  3.3 ,  3.  ,  3.3 ,  3.3 ,  4.3 ],
           [ 1.3 ,  2.99,  3.7 ,  3.  ,  3.3 ,  3.3 ,  4.3 ],
           [ 1.3 ,  2.3 ,  3.7 ,  3.  ,  3.3 ,  3.3 ,  4.3 ]])]
    

提交回复
热议问题