Numpy Argmax, Explained

This tutorial explains how to use the Numpy argmax function. It explains the syntax of np.argmax, and also shows step-by-step examples.

You can click on any of the links below, and it will take you to the appropriate section of the tutorial.

Table of Contents:

• Introduction to Numpy Argmax
• Syntax
• Examples

Having said that, if you’re new to Numpy, you should probably read the whole tutorial. It will make more sense if you read from start to finish.

A quick introduction to Numpy Argmax

Let’s start off with a quick introduction to the argmax function.

The Numpy argmax function often confuses people, but it starts to make sense once someone explains it clearly (which I’m going to try to do).

Essentially, the argmax function returns the index of the maximum value of a Numpy array.

It’s somewhat similar to the Numpy maximum function, but instead of returning the maximum value, it returns the index of the maximum value.

To really explain that, I’m going to quickly review some Numpy and Python basics. This will hopefully make it easier to understand.

Numpy arrays store data

First, let’s quickly review what a Numpy array is.

A Numpy array is a data structure that stores data in a grid format. Although there are exceptions, Numpy arrays almost always store numeric data.

So for example, in the simple Numpy array above, we have 5 values, arranged in a 1 dimensional array.

The Numpy array is essentially a grid-like data structure that stores numeric data.

Locations in a Numpy array have an “index”

The next thing you need to know is that every location in a Numpy array has a position.

It gets a little more complicated for 2D arrays, so let’s keep things simple and look again at a 1D array.

If we have a 1-dimensional array, every location in that array has a sort of address. In Python, we call that address the “index”.

An index for a Numpy array works almost exactly the same as the index for other Python objects. Many other Python data structures – like lists and tuples – use indexes.

Just like the indexes for those structures, Numpy array indexes start at 0. Additionally, we can use those index values to identify or retrieve specific elements of an array.

Numpy Argmax Identifies the Maximum Value and Returns the Associated Index

Now, let’s bring this back to the argmax function.

When we use Numpy argmax, the function identifies the maximum value in the array.

But instead of retrieving the value, Numpy argmax retrieves the index that’s associated with the maximum value.

That’s really it! The np.argmax function simply returns the index of the maximum value in the array.

Having said that, there are some more complicated ways of using the function. For example, you can use the function along particular axes and retrieve the index of the maximum value for a particular array axis.

That’s a little more complicated. I’ll show you how to do that in the examples section, but before I do that, we should look at the syntax first.

The syntax of np.argmax

Ok. Let’s take a look at the syntax.

A quick note

One quick note before we dive in.

This syntax explanation (and the examples below) assume that you’ve imported Numpy with the alias ‘np‘.

You can do that with the code import numpy as np.

This is the common convention among Python data scientists, and we’ll be sticking with it here.

With that said, let’s look at the exact syntax.

np.argmax syntax

The syntax of np.argmax is really pretty simple.

You call the function as np.argmax().

Then, inside of the parenthesis, you have a few parameters that you can use to control exactly how the function works.

Let’s look at those.

The parameters of np.argmax

The np.argmax function really only has 3 parameters:

• a
• axis
• out

The out parameter is somewhat rarely used, so we’re not going to discuss it here.

But let’s quickly look at the a parameter and axis parameter.

a (required)

The a parameter enables you to specify the input array that you want to operate on.

So if you want to operate on an array called myarray, you can call the function as np.argmax(a = myarray).

Keep in mind that you need to provide an argument to this parameter.

Having said that, you don’t need to explicitly use this parameter. Instead, you can pass in an argument by position like this: np.argmax(myarray). The argmax function will assume that the first argument to the function is the input array to be passed to the a= parameter.

Also note that this parameter will accept many data structures as arguments. Typically, we’ll pass in a Numpy array as the argument, but the np.argmax function will also accept “array like” objects, such as Python lists.

axis

The axis parameter enables you to control the axis along which to use argmax.

Remember: Numpy arrays have axes. Axes are like directions along the numpy array.

Keep in mind, that the axis parameter is optional. If you use it, np.argmax will retrieve the index values for the maxima along particular axes. But if you don’t use it, then argmax will flatten out the array and retrieve the index of the maxima of the flattened array.

To be honest, how axes work is little difficult to understand without examples. So I’ll show you some examples in the examples section bellow. I also strongly recommend that you read our tutorial that explains Numpy axes.

Examples of how to use Numpy Argmax

Ok. Let’s take a look at a few examples.

Things almost always make more sense when you can look at some examples, but that’s particularly true with np.argmax.

Examples:

• Use argmax on a 1 dimensional array
• Apply argmax to a 2 dimensional array
• Use argmax along axis 0
• Use argmax along axis 1

Run this code first

Before you run any of the examples, you need to import Numpy.

You can do that with the following code:

import numpy as np

When we do this, we’ll be able to call our Numpy functions starting with the alias ‘np‘.

EXAMPLE 1: Use argmax on a 1 dimensional array

Ok, let’s start simple.

Here, we’re going to identify the index of the maximum value of a 1-dimensional Numpy array.

First, let’s just create the array:

my_1d_array = np.array([1,2,3,100,5])

Next, let’s apply np.argmax.

np.argmax(a = my_1d_array)

OUT:

3

Explanation

This one is pretty simple.

There are several elements in this array. The maximum value of the array is 100.

The maximum value (100) is at index position 3, so argmax returns the value ‘3’.

EXAMPLE 2: Apply argmax to a 2 dimensional array

Let’s take a look at a slightly more complicated example.

Let’s look at how argmax works with a 2-dimensional array.

First, let’s just create our array with the np.array function.

my_2d_array = np.array([[100,2,3],[4,5,600]])

Notice the large values 100 and 600 in the array.

Ok. Now, let’s apply np.argmax.

np.argmax(my_2d_array)

OUT:

5

Explanation

Here the output is 5. Why?

Here, we’re operating on a 2-dimensional array.

By default, if we’re working with a 2D array and we do not specify an axis, the Numpy argmax function will apply a 2-step process.

First, it will flatten out the array to a 1-dimensional array. (Note, it does this for 2D arrays but also for higher dimensional arrays).

Second, it applies the argmax function to the flattened array.

In this case, when we flatten out the array, the maximum value, 600, is at index position 5 of the flattened array.

So, numpy.argmax returns the value 5 in this case.

EXAMPLE 3: Use argmax along axis 0

Now, let’s apply argmax to a 2D array, and also use the axis parameter.

Here, we’re going to set axis = 0.

In this example, we’ll re-use the array that we created in example 2, but here’s the code to recreate it, in case you didn’t run example 2.

my_2d_array = np.array([[100,2,3],[4,5,600]])

Next, let’s apply Numpy argmax with axis = 0:

np.argmax(my_2d_array, axis = 0)

OUT:

array([0, 1, 1])

Explanation

This is a little more complicated, and it’s harder to understand, but let’s break it down.

Remember: Numpy axes are like directions along a Numpy array. For a 2D array, the axis-0 direction points downward against the rows.

When we apply Numpy argmax in the axis-0 direction, it identifies the maximum along that axis and returns the index.

This still might confuse people, so let’s look carefully.

When we set axis = 0, we’re applying argmax in the axis-0 direction, which is downward here.

Along axis-0, every row has an index, so you can see “row 0” and “row 1”. (Remember, all Python indexes start at 0, so the “first” row is actually the 0th row.)

From there, argmax is just looking for the maximum value in the axis 0 direction, and returning the row index.

So 100 is the maximum value in the first column, and the row index of that value is 0.

The maximum value in the second column is 5, which is in row 1.

Similarly, the maximum value in the third column is 600, which is also in row 1.

So the output is the indexes of the maximum values in the axis-0 direction. The output is [0, 1, 1].

Effectively, when we set axis = 0, it’s like applying argmax along the columns.

EXAMPLE 4: Use argmax along axis 1

Ok. Let’s do one more example.

Let’s apply argmax in the axis 1 direction.

First, let’s create our array (the same array as the previous two examples):

my_2d_array = np.array([[100,2,3],[4,5,600]])

And now, let’s apply argmax.

np.argmax(my_2d_array, axis = 1)

OUT:

array([0, 2])

Explanation

This one is also a little hard to understand, and to understand it, you really need to know how Numpy axes work.

Here, we’re applying np.argmax along axis-1. Remember: for 2D Numpy arrays, axis-1 points horizontally across the columns.

So when we set axis = 1, argmax identifies the maximum value for every row. And it returns the column index of that maximum value.

So for the first row, the maximum value is 100. That value has a column index of 0.

For the second row, the maximum value is 600. That value has a column index of 2.

So the output is the column indexes of the maximum values … [0,2].

If you’re still confused, read more about indexes and axes

I’ve tried to show really clear examples here, but I do realize that Numpy argmax is a little hard to wrap your head around.

Ultimately, to understand this function, you really need to understand Numpy indexes. You also really need to understand how axes work … so if you haven’t already, you should read our tutorial that explains Numpy axes.

Leave Your Questions in the Comments Section

If you have any other questions about Numpy argmax, just leave your questions in the comments section near the bottom of the page.

Join our course to learn more about Numpy

In this tutorial, I’ve shown you how to use one Numpy function, Numpy argmax.

Numpy argmax is useful for some tasks, but if you’re working with numeric data in Python, there’s a lot more to learn. You’ll probably have to learn a lot more about Numpy.

If you’re serious about learning Numpy, you should consider joining our premium course called Numpy Mastery.

Numpy Mastery will teach you everything you need to know about Numpy, including:

• How to create Numpy arrays
• How to use the Numpy random functions
• What the “Numpy random seed” function does
• How to Numpy axes work
• How to reshape, split, and combine your Numpy arrays
• Applying mathematical operations on Numpy arrays
• and more …

Additionally, when you join the course, you’ll discover our unique practice system that will enable you to memorize all of the syntax that you learn.

As long as you practice like we show you, you’ll master all of the critical Numpy syntax within a few weeks.

If you have trouble remembering Numpy syntax, this is the course you’ve been looking for.

Find out more here:

Learn More About Numpy Mastery