在numpy中实现轴式赢家通吃的pythonic方法

Question

我想知道什么是最简洁和最典型的方法，将2Dnumpy数组的每一行中只保留最大元素，而将所有其他元素设置为零。示例：

给定下面的numpy数组：

a = [ [1, 8, 3 ,6],
      [5, 5, 60, 1],
      [63,9, 9, 23] ]

我希望答案是：

b = [ [0, 8, 0, 0],
      [0, 0, 60, 0],
      [63,0, 0, 0 ] ]

我可以想出几种方法来解决这个问题，但我感兴趣的是是否有python函数来快速解决这个问题。

提前谢谢你

Answer 1

您可以使用np.max沿一个轴获取最大值，然后使用np.where将非最大值元素清零：

np.where(a == a.max(axis=1, keepdims=True), a, 0)

keepdims=True参数在取最大值后保留单例维度(即a.max(1, keepdims=True).shape == (3, 1))，这简化了针对a的广播。

Answer 2

我不知道python是什么，所以我假设大多数python特定的语法都是python。它使用了两个列表理解，这是python的特性。但在这种情况下，它可能不会那么简洁。

b = [[y if y == max(x) else 0 for y in x] for x in a ]