Numpy 创建一个掩码矩阵，其中在最大值索引处为真

给定一个矩阵

a=[[1,5,3]，[4,2,6]]

我想创建一个掩码矩阵，其中

1

值位于列最大值索引

mask=[[0,1,0]，[1,0,1]

。如何在numpy中执行此操作而不改变任何创建的数组。

您可以尝试：

# if A is not an np array already
A = np.array(A)

(A == np.amax(A, keepdims=True, axis=0)).astype(int)

输出：

array([[0, 1, 0],
       [1, 0, 1]])

---

另一种类似的方法是将0添加到布尔矩阵以将其转换为数字：

mask = (A == np.amax(A,keepdims=True,axis=0)) + 0

[[0 1 0]
 [1 0 1]]

对于较大的矩阵，这个可能更快：

mask = np.zeros_like(A)
mask[np.argmax(A,axis=0),np.arange(A.shape[1])] = 1