如何获得矢量化的 np.argmax 随机决胜局？ [复制]答案

【问题标题】：How to get a vectorized np.argmax random tiebreaker? [duplicate]如何获得矢量化的 np.argmax 随机决胜局？ [复制]
【发布时间】：2019-05-30 06:59:58
【问题描述】：

我知道我可以通过输入一个二维数组并指定一个轴来矢量化 np.argmax，例如：np.argmax(2Darray,axis=1) 以获得每行的最大索引。

我知道如果两个条目在单个 1D 向量中相等，我希望返回最大索引，我可以通过 np.random.choice(np.flatnonzero(1Dvector == 1Dvector.max()) )

问题是，我怎样才能同时做到这两点？即：如何对 np.argmax 进行矢量化，从而使相等的条目随机平局？

【问题讨论】：

你检查了吗：stackoverflow.com/questions/17568612/…
看起来这仅适用于一维向量。我找不到 np.argwhere 的轴参数。此外，我想返回最大值而不是获取最大索引列表，尽管我确信如果 np.argwhere 可以被矢量化，那部分将是微不足道的。

标签： numpy

【解决方案1】：

这是一种方法。对于大数据，可以考虑用更便宜的东西替换permutation。我已经硬编码了axis=1，但这不应该掩盖原则。

def fair_argmax_2D(a):
    y, x = np.where((a.T==a.max(1)).T)
    aux = np.random.permutation(len(y))
    xa = np.empty_like(x)
    xa[aux] = x
    return xa[np.maximum.reduceat(aux, np.where(np.diff(y, prepend=-1))[0])]

a = np.random.randint(0,5,(4,5))
a
# array([[2, 2, 2, 2, 1],
#        [3, 3, 3, 3, 2],
#        [3, 4, 2, 1, 4],
#        [3, 2, 4, 2, 1]])

# draw 10000 times
res = np.array([fair_argmax_2D(a) for _ in range(10000)])

# check
np.array([np.bincount(r, None, 5) for r in res.T])
# array([[ 2447,  2567,  2449,  2537,     0],
#        [ 2511,  2465,  2536,  2488,     0],
#        [    0,  5048,     0,     0,  4952],
#        [    0,     0, 10000,     0,     0]])

【讨论】：

很棒的解决方案，谢谢！