How to move identical elements into a numpy array in subarrays

Question

How to efficiently move identical elements from a sorted numpy array to subarrays?

from here:

import numpy as np a=np.array([0,0,1,1,1,3,5,5,5]) 

here:

 a2=array([[0, 0], [1, 1, 1], [3], [5, 5, 5]], dtype=object) 

Answer 1

One approach would be to get places of shifts where numbers change and use these indices to split the input array into subarrays. To find these indices, you can use np.nonzero for a differentiated array, and then use np.split to split, for example:

 np.split(a,np.nonzero(np.diff(a))[0]+1) 

Run Example -

 In [42]: a Out[42]: array([2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 5, 5, 6, 6, 6]) In [43]: np.split(a,np.nonzero(np.diff(a))[0]+1) Out[43]: [array([2, 2, 2, 2]), array([3, 3, 3, 3]), array([4, 4, 4, 4, 4, 4, 4]), array([5, 5]), array([6, 6, 6])] 

Answer 2

One way to do this is to use itertools.groupby . Example -

 result = np.array([list(g) for _,g in groupby(a)]) 

This will work for regular sorted lists as well, not just numpy arrays.

Demo -

 In [24]: import numpy as np In [25]: a=np.array([0,0,1,1,1,3,5,5,5]) In [26]: from itertools import groupby In [27]: result = np.array([list(g) for _,g in groupby(a)]) In [28]: result Out[28]: array([[0, 0], [1, 1, 1], [3], [5, 5, 5]], dtype=object) 

---

Comparing time with a different approach -

 In [29]: %timeit np.array([list(g) for _,g in groupby(a)]) The slowest run took 6.10 times longer than the fastest. This could mean that an intermediate result is being cached 100000 loops, best of 3: 9.86 µs per loop In [30]: %timeit np.split(a,np.where(np.diff(a)>0)[0]+1) 10000 loops, best of 3: 29.2 µs per loop In [31]: %timeit np.array([list(g) for _,g in groupby(a)]) 100000 loops, best of 3: 10.5 µs per loop In [33]: %timeit np.split(a,np.nonzero(np.diff(a))[0]+1) The slowest run took 4.32 times longer than the fastest. This could mean that an intermediate result is being cached 10000 loops, best of 3: 25.2 µs per loop