Argmax区分具有相同值的数组和具有最大第零个分量的数组

Question

我正在实现argmax函数formnumpy库，以获取向量中最大元素的索引。下面是我的代码

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

从上面的代码中，仅通过查看idx变量，就无法区分argmax的输入是a还是b

我的代码有一个类似的设置，我不知道输入数据的先验性，并且可以获取我收到的数据的argmax。数据可以包含所有相同整数值的条目，也可以在五个索引位置中的任意位置有一个最大值（与其他值相比）。但是，当最大值出现在第0个位置时，我的其余代码将失败，因为我假设第一个索引位置包含最大值

python中是否有任何其他形式的argmax可以返回None（或除0以外的任何形式），以指示输入数据中没有一个最大值

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