Pythonic 方式在 Numpy 中实现轴向赢家通吃

Question

我想知道在 2D numpy 数组的每一行中仅保留最大元素，同时将所有其他元素设置为零的最简洁和 pythonic 方法是什么。例：

给定以下 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

不知道什么是pythonic，所以我假设大多数python特定语法的方式是pythonic。它使用了两个列表理解，这是python的功能。但这样可能就不那么简洁了。

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