如何将numpy数组中的相同元素移动到子数组中

Question

How do I efficiently move identical elements from a sorted numpy array into subarrays?

from here:

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

to here:

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

Answer 1

One approach would be to get the places of shifts, where the numbers change and use those indices to split the input array into subarrays. For finding those indices, you can use np.nonzero on a differentiated array and then use np.split for splitting, like so -

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

Sample run -

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 method to do this would be using itertools.groupby . Example -

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

This would work for normal 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)

---

Timing comparison with the other 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