numpy 中轴向赢家通吃的 pythonic 方式

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 的广播。