创建一个遮罩矩阵，其最大索引值为true

Question

给出一个矩阵A = [[1, 5, 3], [4,2,6]]，我想创建一个掩码矩阵，其中1的值位于列最大值索引mask = [[0, 1, 0], [1, 0, 1]]处。如何在numpy中执行此操作而不更改任何创建的数组。

Answer 1

您可以尝试：

# 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]])

Answer 2

另一种类似的方法是将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