numpy中的argmax()方法只返回第一个最大值，如何找到所有最大值的索引

Question

我有下面的代码

import numpy as np
np.argmax(np.array([5,5,4]))

这只返回0。我期望它将返回出现最大值的所有索引，即0&1

有什么办法做到这一点吗？

Answer 1

您可以使用numpy.where来查找数组最大值的所有索引。

import numpy as np
a = np.array([5,5,4])
res = np.where(a == a.max())[0]
print(res)
# [0 1]