Python Programming

Splitting an array

You can split an array evenly into multiple sub-arrays with np.split(), giving the number of splits and also the axis as input arguments.

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

## split on axis 1 (columns) into 3 even-sized sub-arrays
y = np.split(x, 3, axis=1)  
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]]

## split on axis 0 (rows) into 3 even-sized sub-arrays
y = np.split(x, 3, axis=0)
print(y[0])
## [[1 2 3 4 5 6 7 8 9]]

print(y[1])
## [[10 11 12 13 14 15 16 17 18]]

print(y[2])
## []

A similar function np.array_split() also allows you to split an array without needing to be strictly even. For example, you can divide 10 columns into 3 sub-arrays.

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

## split on axis 1 (columns) at columns 3 and 5.
y = np.split(x, [3, 5], axis=1)  
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]]

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]