NumPy ndarray.partition()

The numpy.ndarray.partition() method rearranges the elements of an array so that the element at the given index kth is in its correct sorted position. All smaller elements appear before kth, and all larger elements appear after it (though not necessarily in sorted order).

Syntax

</>
Copy
ndarray.partition(kth, axis=-1, kind='introselect', order=None)

Parameters

ParameterTypeDescription
kthint or sequence of intsElement index (or multiple indices) to partition by. The element(s) at this position will be placed in their final sorted location.
axisint, optionalAxis along which partitioning is applied. Default is -1 (last axis).
kind{‘introselect’}, optionalSelection algorithm used. Currently, only 'introselect' is supported.
orderstr or list of str, optionalFor structured arrays, specifies the field(s) to sort by.

Return Value

Rearranges the array in-place such that the kth element is in its final sorted position. It does not return a new array but modifies the original array.


Examples

1. Partitioning a 1D Array

In this example, we partition a 1D array based on a given index.

</>
Copy
import numpy as np  # Import NumPy library

# Define a 1D NumPy array
arr = np.array([10, 5, 8, 3, 1, 6, 9])

# Partition the array such that the 3rd element (index 3) is in its correct position
arr.partition(3)

print("Partitioned array:", arr)

Output:

Partitioned array: [1 3 5 6 9 8 10]

Here, the element at index 3 is placed in its correct sorted position, with smaller elements on the left and larger elements on the right.

2. Partitioning Along a Specific Axis

We apply partition() on a 2D array along a specified axis.

</>
Copy
import numpy as np 

# Define a 2D NumPy array
arr = np.array([[7, 2, 1], 
                [8, 6, 4], 
                [3, 9, 5]])

# Partition along columns (axis=0) such that the second element (index 1) in each column is in place
arr.partition(1, axis=0)

print("Partitioned array:\n", arr)

Output:

Partitioned array:
[[3 2 1]
 [7 6 4]
 [8 9 5]]

Each column is rearranged so that the element at index 1 is in its correct position.

3. Partitioning with Multiple Indices

We partition using multiple kth indices.

</>
Copy
import numpy as np 

# Define a 1D array
arr = np.array([15, 3, 20, 10, 5, 7])

# Partition such that elements at indices 1 and 3 are correctly positioned
arr.partition([1, 3])

print("Partitioned array:", arr)

Output:

Partitioned array: [ 3  5  7 10 15 20]

The elements at indices 1 and 3 are placed in their correct sorted positions.

4. Partitioning a Structured Array

We partition a structured array using the order parameter.

</>
Copy
import numpy as np 

# Create a structured array with 'name' and 'age' fields
data = np.array([('Arjun', 30), ('Bhairav', 25), ('Charlie', 35)], 
                dtype=[('name', 'U10'), ('age', 'i4')])

# Partition based on the 'age' field
data.partition(1, order='age')

print("Partitioned array:", data)

Output:

Partitioned array: [('Bhairav', 25) ('Arjun', 30) ('Charlie', 35)]

The order='age' argument ensures that the partitioning happens based on the age field, placing the element at index 1 in its correct position.