NumPy array_split() - Split an array to subarrays

A NumPy array can be split by array_split().

import numpy

a = numpy.array([1, 2, 3, 4, 5, 6])

b = numpy.array_split(a, 2)
c = numpy.array_split(a, 3)

print(b)  # [array([1, 2, 3]), array([4, 5, 6])]
print(c)  # [array([1, 2]), array([3, 4]), array([5, 6])]

The first argument is the array you want to split and the second one is the number of splitting. If this number is not a divisor of the number of elements in the original array, the split arrays doesn't have the same dimension.

import numpy

a = numpy.array([1, 2, 3, 4, 5, 6])

b = numpy.array_split(a, 4)
c = numpy.array_split(a, 5)

print(b)  # [array([1, 2]), array([3, 4]), array([5]), array([6])]
print(c)  # [array([1, 2]), array([3]), array([4]), array([5]), array([6])]

If the splitting count is 0, Python raises a ValueError exception. And if this value is more than the number of elements in the original array, the result has extra empty arrays as follows.

import numpy

a = numpy.array([1, 2, 3, 4, 5, 6])

b = numpy.array_split(a, 1)
c = numpy.array_split(a, 6)

print(b)  # [array([1, 2, 3, 4, 5, 6])]
print(c)  # [array([1]), array([2]), array([3]), array([4]), array([5]), array([6])]

d = numpy.array_split(a, 0)
# ValueError: number sections must be larger than 0.

e = numpy.array_split(a, 7)
print(e)  # [array([1]), array([2]), array([3]), array([4]), array([5]), array([6]), array([], dtype=int64)]

Note: Split a Python list using list comprehension

Side note, a Python list can be split to sublists (chunks) by list comprehension

a = [1, 2, 3, 4, 5, 6]

b = [a[x:x + 2] for x in range(0, len(a), 2)]
c = [a[x:x + 3] for x in range(0, len(a), 3)]

print(b)  # [[1, 2], [3, 4], [5, 6]]
print(c)  # [[1, 2, 3], [4, 5, 6]]