获取具有随机平局的二维数组的行最大值的列索引答案

【问题标题】：Get column indices of row-wise maximum values of a 2D array with random tie-breaking获取具有随机平局的二维数组的行最大值的列索引
【发布时间】：2023-01-12 09:04:47
【问题描述】：

给定一个 2D numpy 数组，我想用每行最大值的列索引构造一个数组。到目前为止，arr.argmax(1) 运行良好。但是，对于我的具体情况，对于某些行，2 列或更多列可能包含最大值。在那种情况下，我想随机选择一个列索引（而不是第一个索引，因为它是 .argmax(1) 的情况）。

例如，对于以下arr：

arr = np.array([
    [0, 1, 0],
    [1, 1, 0],
    [2, 1, 3],
    [3, 2, 2]
])

可以有两种可能的结果：array([1, 0, 2, 0]) 和 array([1, 1, 2, 0])，每一种都以 1/2 的概率被选中。

我使用列表理解编写了返回预期输出的代码：

idx = np.arange(arr.shape[1])
ans = [np.random.choice(idx[ix]) for ix in arr == arr.max(1, keepdims=True)]

但我正在寻找优化的 numpy 解决方案。换句话说，我如何用 numpy 方法替换列表理解以使代码适用于更大的数组？

【问题讨论】：

标签： python numpy

【解决方案1】：

使用 scipy.stats.rankdata 和 apply_along_axis 如下。

import numpy as np
from scipy.stats import rankdata
ranks = rankdata(-arr, axis = 1, method = "min")
func = lambda x: np.random.choice(np.where(x==1)[0])
idx = np.apply_along_axis(func, 1, ranks)

print(idx)

它返回 [1 0 2 0] 或 [1 1 2 0]。

主要思想是rankdata计算每一行中每个值的排名，最大值为1。func随机选择一个对应值为1的索引。最后，apply_along_axis将func应用于每个arr 行。

【讨论】：