Partition array into N chunks with Numpy Partition array into N chunks with Numpy python python

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]])]


I believe that you're looking for numpy.split or possibly numpy.array_split if the number of sections doesn't need to divide the size of the array properly.