Understanding numpy.array_split() function (4 examples)

Introduction

The NumPy library is an essential tool in the Python ecosystem for efficient manipulation and processing of numerical data. Among its vast array of functionalities, the array_split() function is a versatile method for splitting arrays into multiple sub-arrays. This tutorial aims to provide a comprehensive understanding of how to use the numpy.array_split() function through practical examples.

What is numpy.array_split()?

The numpy.array_split() function splits an array into multiple sub-arrays as specified by the user. Unlike split(), array_split() allows for non-uniform divisions, meaning the sub-arrays can have unequal sizes. This functionality is particularly useful when handling datasets that cannot be equally divided.

Syntax:

numpy.array_split(ary, indices_or_sections, axis=0)

Parameters:

• ary: array_like. The input array to be split.
• indices_or_sections: int or 1-D array_like. If an integer, N, the array will be divided into N sub-arrays along the specified axis. If an array, the elements represent the points at which the array is split along the specified axis.
• axis: int, optional. The axis along which to split the array. Default is 0.

Returns:

• sub_arrays: list of ndarrays. A list of sub-arrays into which the input array has been split.

Example #1 – Basic Usage

Let’s start with a simple example to demonstrate the basic usage of array_split():

import numpy as np

# Create an array
arr = np.arange(10)
# Split the array into 3 sub-arrays
result = np.array_split(arr, 3)

print(result)

Output:

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

As you can see, the original array was split into three sub-arrays, with the first sub-array being slightly larger to accommodate the non-uniform division.

Example #2 – Split Along Different Axes

NumPy arrays can have multiple dimensions, and array_split() can split arrays along any given axis. Let’s look at splitting a 2D array:

import numpy as np

# Create a 2D array
arr2d = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# Split the array vertically into 2 parts
result = np.array_split(arr2d, 2, axis=1)

print(result)

Output:

 [array([[1, 2], [4, 5], [7, 8]]), array([[3], [6], [9]])] 

This example demonstrates splitting a 2D array into two sub-arrays vertically. Notice how array_split() allows for a non-uniform split along the specified axis.

Example #3 – Using Advanced Indexing

Advanced indexing can further customize how arrays are split. Imagine you want more control over the indices where the split occurs:

import numpy as np

# Create an array
arr = np.array([10, 20, 30, 40, 50, 60, 70, 80, 90, 100])
# Split the array at specified indices
result = np.array_split(arr, [3, 5, 6])

print(result)

Output:

 [array([10, 20, 30]), array([40, 50]), array([60]), array([70, 80, 90, 100])] 

Here, array_split() divided the array into four sub-arrays at the specified indices. This level of control allows for fine-tuned manipulation of data.

Example #4 – Handling Multidimensional Arrays

Finally, let’s delve into splitting high-dimensional arrays. For applications involving multidimensional datasets, like image data or 3D models, array_split() is invaluable:

import numpy as np

# Create a 3D array
arr3d = np.random.randint(1, 100, size=(4,4,4))
# Split the 3D array into 2 parts along the third axis (depth)
result = np.array_split(arr3d, 2, axis=2)

print('First sub-array shape:', result[0].shape)
print('Second sub-array shape:', result[1].shape)

Output:

First sub-array shape: (4, 4, 2)
Second sub-array shape: (4, 4, 2)

This operation splits a 3D array into two parts along its depth, demonstrating the function’s flexibility across different dimensions and the practical applications for complex datasets.

Conclusion

The numpy.array_split() function is a powerful and flexible tool for array manipulation, allowing for non-uniform divisions and customizable splits across any axis. Through these examples, ranging from basic to advanced, we’ve explored its vast potential in handling both simple and complex data structures. With this knowledge, you’re now better equipped to manipulate arrays efficiently in your data processing tasks.