python - 来自 numpy 数组中的最大值的掩码，特定轴

Question

输入示例:

我有一个 numpy 数组，例如

a=np.array([[0,1], [2, 1], [4, 8]])

期望的输出:

我想生成一个掩码数组，其最大值沿给定轴(在我的情况下为轴 1)为 True，所有其他轴为 False。例如在这种情况下

mask = np.array([[False, True], [True, False], [False, True]])

尝试:

我尝试过使用 np.amax 的方法，但这会返回展平列表中的最大值:

>>> np.amax(a, axis=1)
array([1, 2, 8])

和 np.argmax 类似地返回沿该轴的最大值的索引。

>>> np.argmax(a, axis=1)
array([1, 0, 1])

我可以以某种方式对此进行迭代，但是一旦这些数组变大，我希望解决方案在 numpy 中保持原生。

Answer 1

方法#1

使用 broadcasting ，我们可以使用与最大值的比较，同时保持暗淡以促进广播 -

a.max(axis=1,keepdims=1) == a

sample 运行-

In [83]: a
Out[83]: 
array([[0, 1],
       [2, 1],
       [4, 8]])

In [84]: a.max(axis=1,keepdims=1) == a
Out[84]: 
array([[False,  True],
       [ True, False],
       [False,  True]], dtype=bool)

方法#2

对于另一种情况，argmax 索引对于沿列的索引范围的 broadcasted-comparison 或者 -

In [92]: a.argmax(axis=1)[:,None] == range(a.shape[1])
Out[92]: 
array([[False,  True],
       [ True, False],
       [False,  True]], dtype=bool)

方法#3

要完成设置，如果我们正在寻找性能，请使用初始化，然后使用 advanced-indexing -

out = np.zeros(a.shape, dtype=bool)
out[np.arange(len(a)), a.argmax(axis=1)] = 1