时间戳

首页 > 解决方案 > 如何获得向量。从使用 argmax 的二维 numpy 数组？

python - 如何获得向量。从使用 argmax 的二维 numpy 数组？

问题描述

我有以下内容numpy.ndarray：

测试：

array([[  2.55053788e-01,   6.25406146e-01,   1.19271643e-01,
          2.68359261e-04],
       [  2.59611636e-01,   0.19562805e-01,   1.20518960e-01,
          3.06535745e-01],
       [  8.52524495e-01,   5.24317825e-01,   1.22851081e-01,
          3.06610862e-04],
       [  2.55068243e-01,   6.24345124e-01,   1.20263465e-01,
          3.23178538e-04],
       [  2.46678621e-01,   6.29301071e-01,   1.23693809e-01,
          3.26490292e-04]], dtype=float32)

如果我这样做testVec = np.argmax(testArr)，我会得到一个号码。如何获得 0、1 或 2 的向量，具体取决于 2D 数组每行中的最大值testArr？

预期输出：

[1, 3, 0, 1, 1]

标签： pythonnumpyargmax

解决方案

如果您documentation看一下，您会看到有一个axis参数允许您选择要沿哪个轴执行操作。从文档：

返回沿轴的最大值的索引。

在这种情况下，您想要：

np.argmax(a, axis=1)
# array([1, 3, 0, 1, 1], dtype=int64)

推荐阅读