找出矩阵在行轴上的最大值，将其他值替换为零

Question

A = [[2,2,4,2,2,2]
[2,6,2,2,2,2]
[2,2,2,2,8,2]]

我想让矩阵B等于

B = [[0,0,4,0,0,0]
[0,6,0,0,0,0]
[0,0,0,0,8,0]]

所以我想找到每一行的最大值并将其他值替换为0。有没有办法做到这一点，而不使用for循环?提前感谢您的意见。

Answer 1

您可以直接获取每行的最大值，而不是查看argmax，然后屏蔽较低的元素并将其替换为零:

在适当的位置看起来像(这里True代表keepdims=True):

>>> A[A < A.max(1, True)] = 0
>>> A
array([[0, 0, 4, 0, 0, 0],
[0, 6, 0, 0, 0, 0],
[0, 0, 0, 0, 8, 0]])

一个不合适的选择是使用np.where:

>>> np.where(A == A.max(1, True), A, 0)
array([[0, 0, 4, 0, 0, 0],
[0, 6, 0, 0, 0, 0],
[0, 0, 0, 0, 8, 0]])