# np.argpartition() — A Simple Illustrated Guide

In Python, the `numpy.argpartition()` function returns the indices that would partition an array along with a given axis based on the specified kth element(s).

All elements smaller than the kth element will be moved before it and all larger elements behind it. The element order in the partitions is undefined.

If provided with a sequence of kth, it will partition all of them into their sorted position at once.

Here is the argument table of `np.argpartition()`.

If it sounds great to you, please continue reading, and you will fully understand the `np.argpartition()` function through Python code snippets and vivid visualization.

This tutorial is about `np.argpartition()` function.

• Concretely, I will introduce its syntax and arguments.
• Then, you will learn some basic examples of this function.
• Finally, I will address three top questions about `np.argpartition()`, including `np.argpartition` 2d array/axis, np.argpartition order, and `np.argpartition` ignore `np.nan`.

You can find all codes in this tutorial here.

## Syntax and Arguments

Here is the syntax of `np.argpartition()`:

```# Syntax
numpy.argpartition(a, kth[, axis=- 1[,
kind='introselect'[, order=None]]])```

Here is the argument table of `np.argpartition()`:

Don’ panic! I wil explain the kth, axis, and order argument clearly later!

The output of `np.argpartition()` function is an `index_array`. If the input array `a` is one-dimensional, `a[index_array]` yields a partitioned `a`. Generally, `np.take_along_axis(a, index_array, axis=a)` always yields the partitioned `a`, irrespective of dimensionality.

## Basic Examples

Here is a one-dimensional array code example:

```# Basic Example
import numpy as np
one_dim = np.array([2, 3, 1, 5, 4])
# kth=0 -> partition based on 2(zero index).
partitioned = np.argpartition(one_dim, 0)
print(f'Unpartitioned array: {one_dim}')
print(f'Partitioned array index: {partitioned}')
print(f'Partitioned array: {one_dim[partitioned]}')```

Output:

As you can see, the partitioned array does not guarantee the element order in the partitions. It just separates elements into groups based on the kth argument!

This is a way to do partial sort and has an efficient time complexity O(n).

If you are more interested in sorting array as a whole, you might want to check out the `numpy.argsort()` function and here is my tutorial on numpy.argsort().

## np.argpartition() 2d array/axis

In this part, I will show you how to deploy the argument axis with 2d array. Of course, you can extrapolate this into an n-dimensional array!

Just for your quick reference, here is the argument table of `np.argpartition()`:

Here is the 2d array code example:

```import numpy as np

# axis = 0 [partial sort along the axis = 0. In this case, row-like.]
print('axis = 0')
two_dim = np.array([[1, 4, 3], [3, 2, 1]])
partitioned = np.argpartition(two_dim, kth=0, axis=0)
print(f'Unpartitioned array: {two_dim}')
print(f'Partitioned array index: {partitioned}')
print(f'Partitioned array: {np.take_along_axis(two_dim, partitioned, axis=0)}')

# axis = 1 [partial sort along the axis = 1. In this case, column-like.]
print('-'*85)
print('axis = 1')
two_dim = np.array([[1, 4, 3], [3, 2, 1]])
partitioned = np.argpartition(two_dim, kth=0, axis=1)
print(f'Unpartitioned array: {two_dim}')
print(f'Partitioned array index: {partitioned}')
print(f'Partitioned array: {np.take_along_axis(two_dim, partitioned, axis=1)}')
```

Output:

## np.argpartition() order

In this part, I will show you another exciting argument, order.

Just for your quick reference, here is the argument table of `np.argpartition()`:

Here is a one-dimensional array code example:

```import numpy as np

# order: x -> y
print('order: x -> y')
one_dim = np.array([(1, 3), (2, 1), (1, 1), (2, 3)],
dtype=[('x', float), ('y', float)])
partitioned = np.argpartition(one_dim, kth=0, order=['x', 'y'])
# In this case, same as:
# partitioned = np.argpartition(one_dim, kth=0)
print(f'Unpartitioned array: {one_dim}')
print(f'Partitioned array index: {partitioned}')
print(f'Partitioned array: {one_dim[partitioned]}')

# order: y -> x
print('-'*85)
print('order: y -> x')
one_dim = np.array([(1, 3), (2, 1), (1, 1), (2, 3)],
dtype=[('x', float), ('y', float)])
partitioned = np.argpartition(one_dim, kth=0, order=['y', 'x'])
print(f'Unpartitioned array: {one_dim}')
print(f'Partitioned array index: {partitioned}')
print(f'Partitioned array: {one_dim[partitioned]}')
```

Output:

## np.argpartition() ignore np.nan

When we encounter tricky some `np.nan` in a sorting problem, we can exclude them by counting the number of `np.nan` and leaving them out through indexing.

Here is a case of getting the N biggest element and index in an array containing `np.nan`:

```import numpy as np

# * Do not handle np.nan problem:
print('Do not handle np.nan problem:')
one_dim = np.array([2, 3, 1, np.nan, 5, 4, np.nan])
# kth=0 -> partition based on 2(zero index).
partitioned = np.argpartition(one_dim, 0)
print(f'Unpartitioned array: {one_dim}')
print(f'Partitioned array index: {partitioned}')
print(f'Partitioned array: {one_dim[partitioned]}')

# * Handle np.nan problem:
# np.nan is taken as an infinite large number in Python.
print('-'*85)
print('Handle np.nan problem:')
one_dim = np.array([2, 3, 1, np.nan, 5, 4, np.nan])
# If we want to get the N biggest element and index in an array:
print('(If we want to get the N biggest element and index in an array)')
N = 2
# get the number of np.nan
c = np.isnan(one_dim).sum()
partitioned_idx = np.argpartition(one_dim, -N-c)[-N-c:-c]
partitioned_val = one_dim[partitioned_idx]
print(f'Unpartitioned array: {one_dim}')
print(f'Partitioned array index: {partitioned_idx}')
print(f'Partitioned array: {partitioned_val}')
```

Output:

## Summary

That’s it for our `np.argpartition()` article.

We learned about its syntax, arguments, and basic examples.

We also worked on the top three questions about the `np.argpartition()` function, ranging from `np.argpartition` 2d array/axis, `np.argpartition` order, and `np.argpartition` ignore `np.nan`

By the way, if you want to do a direct partial sort, please refer to the `numpy.partition()` function. I believe that after reading this tutorial, you can understand `numpy.partition()` easily!

Hope you enjoy all this and happy coding!