使用索引数组在 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

setting values in numpy matrix with array of indices