Understanding argmax

It took me a while to figure this function out. Basically argmax returns you the index of the maximum value in the array. Now the array can be 1 dimensional or multiple dimensions. Following are some examples.

1 dimensional

a = [[1,2,3,4,5]]
np.argmax(a)
>>4

The array is 1 dimensional so the function simply returns the index of the maximum value(5) in the array, which is 4.

Multiple dimensions

a = [[1,2,3],[4,5,6]]
np.argmax(a)
>>5

In this example the array is 2 dimensional, with shape (2,3). Since no axis parameter is specified in the function, the numpy library flattens the array to a 1 dimensional array and then returns the index of the maximum value. In this case the array is transformed to [[1,2,3,4,5,6]] and then returns the index of 6, which is 5.

When parameter is axis = 0

a = [[1,2,3],[4,5,6]]
np.argmax(a, axis=0)
>>array([1, 1, 1])

The result here was a bit confusing to me at first. Since the axis is defined to be 0, the function will now try to find the maximum value along the rows of the matrix. The maximum value,6, is in the second row of the matrix. The index of the second row is 1. According to the documentation the dimension specified in the axis parameter will be removed. Since the shape of the original matrix was (2,3) and axis specified as 0, the returned matrix will have a shape of(3,) instead, since the 2 in the original shape(2,3) is removed.The row in which the maximum value was found is now repeated for the same number of elements as the columns in the original matrix i.e. 3.

When parameter is axis = 1

a = [[1,2,3],[4,5,6]]
np.argmax(a, axis=1)
>>array([2, 2])

Same concept as above but now index of the column is returned at which the maximum value is available. In this example the maximum value 6 is in the 3rd column, index 2. The column of the original matrix with shape (2,3) will be removed, transforming to (2,) and so the return array will display two elements, each showing the index of the column in which the maximum value was found.

No argmax returns the position of the largest value. max returns the largest value.

import numpy as np    
A = np.matrix([[1,2,3,33],[4,5,6,66],[7,8,9,99]])

np.argmax(A)  # 11, which is the position of 99

np.argmax(A[:,:])  # 11, which is the position of 99

np.argmax(A[:1])  # 3, which is the position of 33

np.argmax(A[:,2])  # 2, which is the position of 9

np.argmax(A[1:,2])  # 1, which is the position of 9

Here is how argmax works. Let's suppose we have given array

matrix = np.array([[1,2,3],[4,5,6],[7,8,9], [9, 9, 9]])

Now, find the max value from given array

np.max(matrix)

The answer will be -> 9

Now find argmax of given array

np.argmax(matrix)

The answer will be -> 8

How it got 8, let's understand

python will convert array to one dimension, so array will look like

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

so max value is 9 and first occurrence of 9 is at index 8. That's why answer of argmax is 8.

• axis = 0 (column wise max)

Now, find max value column wise

np.argmax(matrix, axis=0)

Index   0  1  2
  0    [1, 2, 3]
  1    [4, 5, 6]
  2    [7, 8, 9]
  3    [9, 9, 9]

• In first column values are 1 4 7 9, max value in first column is 9 and is at index 3
• same for second column, values are 2 5 8 9 and max value in second column is 9 and is at index 3
• for third column values are 3 6 9 9 and max value is 9 and is at index 2 and 3, so first occurrence of 9 is at index 2

so the output will be like [3, 3, 2]

• axis = 1 (row wise)

Now find max value row wise

np.argmax(matrix, axis=1)

    Index   0  1  2
      0    [1, 2, 3]
      1    [4, 5, 6]
      2    [7, 8, 9]
      3    [9, 9, 9]

• for first row values are 1 2 3 and max value is 3 and is at index 2
• for second row values are 4 5 6 and max value is 6 and is at index 2
• for third row values are 7 8 9 and max value is 9 and is at index 2
• for fourth row values are 9 9 9 and max value is 9 and is at index 0 1 2, but first occurrence of 9 is at index 0

so the output will be like [2 2 2 0]