Splitting Arrays

Splitting is the opposite of stacking. It cuts one array into a list of smaller arrays along an axis. Use it to break data into batches or to separate columns.

np.split into equal parts

np.split(arr, n) returns a list of n equal parts. If the length does not divide evenly, it raises ValueError.

import numpy as np

arr = np.arange(6)
print(np.split(arr, 2))

try:
    np.split(np.arange(7), 2)
except ValueError as error:
    print("ValueError:", error)

Output:

[array([0, 1, 2]), array([3, 4, 5])]
ValueError: array split does not result in an equal division

np.array_split for uneven parts

np.array_split allows parts of different sizes. The earlier parts get one extra value each until the remainder is used up.

import numpy as np

for part in np.array_split(np.arange(13), 4):
    print(part)

Output:

[0 1 2 3]
[4 5 6]
[7 8 9]
[10 11 12]

Split at chosen positions

Pass a list instead of a count. Each number is a position where a new part starts.

import numpy as np

for part in np.split(np.arange(10), [3, 7]):
    print(part)

Output:

[0 1 2]
[3 4 5 6]
[7 8 9]

2-D: rows or columns

vsplit cuts between rows and is the same as split(..., axis=0). hsplit cuts between columns and is the same as split(..., axis=1).

import numpy as np

m = np.arange(1, 17).reshape(4, 4)
top, bottom = np.vsplit(m, 2)
print(top)
print(bottom)

left, right = np.hsplit(m, 2)
print(left)
print(right)

Output:

[[1 2 3 4]
 [5 6 7 8]]
[[ 9 10 11 12]
 [13 14 15 16]]
[[ 1  2]
 [ 5  6]
 [ 9 10]
 [13 14]]
[[ 3  4]
 [ 7  8]
 [11 12]
 [15 16]]

3-D: depth

dsplit cuts along the third axis, so it needs an array with at least three dimensions.

import numpy as np

t = np.arange(24).reshape(2, 3, 4)
front, back = np.dsplit(t, 2)
print(t.shape, front.shape, back.shape)
print(front[0])

Output:

(2, 3, 4) (2, 3, 2) (2, 3, 2)
[[0 1]
 [4 5]
 [8 9]]

Key takeaways

• split needs equal parts. array_split allows uneven ones

• A list of positions says where each new part starts

• vsplit cuts rows, hsplit cuts columns, dsplit cuts depth

• The result is a list of arrays

• Next: Iterating Arrays

• Previous: Append, Insert, and Delete