python - 值相等时的 numpy argmax

Question

我有一个 numpy 矩阵，我想获取每一行中最大值的索引。例如

[[1,2,3],[1,3,2],[3,2,1]]

会回来

[0,1,2]

但是，当每行中的最大值超过 1 个时，numpy.argmax 将只返回最小的索引。例如

[[0,0,0],[0,0,0],[0,0,0]]

会回来

[0,0,0]

我可以将默认值(最小索引)更改为其他一些值吗？例如。当最大值相等时，返回1或None，这样上面的结果就是

[1,1,1]
or
[None, None, None]

如果我能在 TensorFlow 中做到这一点，那就更好了。

谢谢!

Answer 1

您可以使用 np.partition 两次找到两个最大值并检查它们是否相等，然后将其用作 np.where 中的掩码来设置默认值:

In [228]: a = np.array([[1, 2, 3, 2], [3, 1, 3, 2], [3, 5, 2, 1]])

In [229]: twomax = np.partition(a, -2)[:, -2:].T

In [230]: default = -1

In [231]: argmax = np.where(twomax[0] != twomax[1], np.argmax(a, -1), default)

In [232]: argmax
Out[232]: array([ 2, -1,  1])