假设我有一个2D numpy数组，表示m个样本在n个类别中的概率（每个样本的概率总和为1）。

假设每个样本只能属于一个类别，我想创建一个与原始数组形状相同的数组，但只包含二进制值，表示哪个类别的概率最高。

示例：

[[0.2, 0.3, 0.5], [0.7, 0.1, 0.1]]

应转换为：

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

看起来amax函数几乎可以实现我想要的功能，但我想得到的是上述描述的指示矩阵，而不是索引。

这看起来很简单，但我无法使用标准的numpy函数实现。我当然可以使用常规的Python循环，但似乎应该有更简单的方法。

如果多个类别的概率相同，我更希望有一个只选择其中一个类别的解决方案（在这种情况下，我不在乎选择哪个类别）。

谢谢！

回答：

这是一种方法：

In [112]: aOut[112]: array([[ 0.2,  0.3,  0.5],       [ 0.7,  0.1,  0.1]])In [113]: a == a.max(axis=1, keepdims=True)Out[113]: array([[False, False,  True],       [ True, False, False]], dtype=bool)In [114]: (a == a.max(axis=1, keepdims=True)).astype(int)Out[114]: array([[0, 0, 1],       [1, 0, 0]])

（但这会为每一行中的最大值的每次出现都标记为True。请查看Divakar的回答，了解如何选择最大值的第一次出现的巧妙方法。）