python - Argmax - differentiate between array having same values and array having largest zeroth component

Question

I am implementing the argmax function form numpy library to get the index of the largest element in a vector. Following is my code.

import numpy as np

a = np.array([2, 0, 0, 0, 0])
idx = np.argmax(a)  # Works fine, returns 0

b = np.array([0, 0, 0, 0, 0])
idx = np.argmax(b)  # Returns 0 again

From the above piece of code, just by looking at idx variable it is impossible to differentiate if the input to argmax was a or b.

My code has a similar setup where I do not know the input data apriori and am allowed to take the argmax of the data I receive. The data can contain entries that are all the same integer values or it can have one maximum value (compared to the rest) in any of the five index positions. When the maximum occurs at the zeroth position, however, the rest of my code fails as I assume that the first index position contains the maximum value.

Is there any other form of argmax in python that can return None (or anything other than 0) to indicate there is no one maximum value in the input data?

question from:https://stackoverflow.com/questions/65944826/argmax-differentiate-between-array-having-same-values-and-array-having-largest

Answer 1

import numpy as np

a = np.array([2, 0, 0, 0, 0])
idx = np.argmax(a) if ~np.all(a == a[0]) else None
print(idx)  # 0

b = np.array([0, 0, 0, 0, 0])
idx = np.argmax(b) if ~np.all(b == b[0]) else None
print(idx)  # None

# Alternative solution

a = np.array([2, 0, 0, 0, 0])
idx = np.argmax(a) - np.all(a == a[0]).astype(int)
print(idx)  # 0

b = np.array([0, 0, 0, 0, 0])
idx = np.argmax(b) - np.all(b == b[0]).astype(int)
print(idx)  # -1