nthy中的轴心胜者全能的pythonic方式

Question

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

给出以下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参数在获取max(即这样a.max(1, keepdims=True).shape == (3, 1))后保持单例维度,这简化了广播a.