在具有索引数组的 numpy 矩阵中设置值

Question

考虑一个 mxn 矩阵 A。我计算了每一行的最大值的索引，产生了一个维度为 m 的数组。

如何使用此索引数组将每行中与 A 形状相同的第二个矩阵 B 中的值设置为 0。

例子：

     A = [[1,2] ,[3,4]] 

     np.argmax(A,1) --> [1,1]

     B = [[1,1] ,[1,1]] 

我希望有：

  B = [[1 0] ,[1 0]] 

我怎样才能做到这一点？

Answer 1

您可以使用：

B[np.arange(B.shape[0]),np.argmax(A,1)] = 0