Partition array into N chunks with Numpy Partition array into N chunks with Numpy numpy numpy

Partition array into N chunks with Numpy


Try numpy.array_split.

From the documentation:

>>> x = np.arange(8.0)>>> np.array_split(x, 3)    [array([ 0.,  1.,  2.]), array([ 3.,  4.,  5.]), array([ 6.,  7.])]

Identical to numpy.split, but won't raise an exception if the groups aren't equal length.

If number of chunks > len(array) you get blank arrays nested inside, to address that - if your split array is saved in a, then you can remove empty arrays by:

[x for x in a if x.size > 0]

Just save that back in a if you wish.


Just some examples on usage of array_split, split, hsplit and vsplit:

n [9]: a = np.random.randint(0,10,[4,4])In [10]: aOut[10]: array([[2, 2, 7, 1],       [5, 0, 3, 1],       [2, 9, 8, 8],       [5, 7, 7, 6]])

Some examples on using array_split:
If you give an array or list as second argument you basically give the indices (before) which to 'cut'

# split rows into 0|1 2|3In [4]: np.array_split(a, [1,3])Out[4]:                                                                                                                       [array([[2, 2, 7, 1]]),                                                                                                        array([[5, 0, 3, 1],                                                                                                                [2, 9, 8, 8]]),                                                                                                         array([[5, 7, 7, 6]])]# split columns into 0| 1 2 3In [5]: np.array_split(a, [1], axis=1)                                                                                           Out[5]:                                                                                                                       [array([[2],                                                                                                                         [5],                                                                                                                          [2],                                                                                                                          [5]]),                                                                                                                  array([[2, 7, 1],                                                                                                                   [0, 3, 1],       [9, 8, 8],       [7, 7, 6]])]

An integer as second arg. specifies the number of equal chunks:

In [6]: np.array_split(a, 2, axis=1)Out[6]: [array([[2, 2],       [5, 0],       [2, 9],       [5, 7]]), array([[7, 1],       [3, 1],       [8, 8],       [7, 6]])]

split works the same but raises an exception if an equal split is not possible

In addition to array_split you can use shortcuts vsplit and hsplit.
vsplit and hsplit are pretty much self-explanatry:

In [11]: np.vsplit(a, 2)Out[11]: [array([[2, 2, 7, 1],       [5, 0, 3, 1]]), array([[2, 9, 8, 8],       [5, 7, 7, 6]])]In [12]: np.hsplit(a, 2)Out[12]: [array([[2, 2],       [5, 0],       [2, 9],       [5, 7]]), array([[7, 1],       [3, 1],       [8, 8],       [7, 6]])]


Not quite an answer, but a long comment with nice formatting of code to the other (correct) answers. If you try the following, you will see that what you are getting are views of the original array, not copies, and that was not the case for the accepted answer in the question you link. Be aware of the possible side effects!

>>> x = np.arange(9.0)>>> a,b,c = np.split(x, 3)>>> aarray([ 0.,  1.,  2.])>>> a[1] = 8>>> aarray([ 0.,  8.,  2.])>>> xarray([ 0.,  8.,  2.,  3.,  4.,  5.,  6.,  7.,  8.])>>> def chunks(l, n):...     """ Yield successive n-sized chunks from l....     """...     for i in xrange(0, len(l), n):...         yield l[i:i+n]... >>> l = range(9)>>> a,b,c = chunks(l, 3)>>> a[0, 1, 2]>>> a[1] = 8>>> a[0, 8, 2]>>> l[0, 1, 2, 3, 4, 5, 6, 7, 8]