Chapter 5: Array manipulation

Array split

face Josiah Wang

You can also do the opposite: you can split an array evenly into multiple sub-arrays with np.split(). You can specify the number of splits and also the axis.

>>> x = np.arange(1, 19).reshape((2, 9))
>>> print(x)
[[ 1  2  3  4  5  6  7  8  9]
 [10 11 12 13 14 15 16 17 18]]
>>> y = np.split(x, 3, axis=1) # split on axis 1 into 3 evenly-sized sub-arrays
>>> print(y[0])
[[ 1  2  3]
 [10 11 12]]
>>> print(y[1])
[[ 4  5  6]
 [13 14 15]]
>>> print(y[2])
[[ 7  8  9]
 [16 17 18]]
>>> y = np.split(x, 2, axis=0) # split on axis 0 into 2 evenly-sized sub-arrays
>>> print(y[0])
[[1 2 3 4 5 6 7 8 9]]
>>> print(y[1])
[[10 11 12 13 14 15 16 17 18]]
>>> y = np.split(x, 4, axis=1) # attempt to split on axis 1 into 4 evenly-sized sub-arrays
ValueError: array split does not result in an equal division

There is a similar function np.array_split() that allows you to split an array without needing to be strictly even. For example, you can divide 10 columns into 3 sub-arrays. You cannot do this with np.split().

You can also specify where to split, by giving a list or tuple as the second argument (instead of an integer).

>>> y = np.split(x, [3, 5], axis=1) # split at columns 3 and 5. 
>>> print(y[0])
[[ 1  2  3]
 [10 11 12]]
>>> print(y[1])
[[ 4  5]
 [13 14]]
>>> print(y[2])
[[ 6  7  8  9]
 [15 16 17 18]]

Like stack, there are also special functions that split at specific axes.

  • np.vsplit(arr, section): same as np.split(arr, section, axis=0) [Vertical split]
  • np.hsplit(arr, section): same as np.split(arr, section, axis=1) [Horizontal split]
  • np.dsplit(arr, section): same as np.split(arr, section, axis=2) [Depth split]