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

时间:2019-02-07 13:20:49

标签: python numpy argmax

我有以下numpy.ndarray

testArr:

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),我会得到一个数字。如何根据二维数组testArr每行中的最大值获得0、1或2的向量?

预期输出:

[1, 3, 0, 1, 1]

2 个答案:

答案 0 :(得分:2)

如果给documentation一眼,就会看到有一个axis参数,该参数使您可以选择要沿哪个轴执行操作。从文档中:

  

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

在这种情况下,您需要:

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

答案 1 :(得分:1)

默认情况下,np.argmax为您提供展平数组中最大值的索引。要获得单个维度上的最大值(例如,每行的最大值),必须指定关键字参数axis。这必须是一个整数:0表示列,1表示行。 (或者如果数组具有n个维,则为n-1之间的任何整数。)

import numpy as np
testArr = np.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=np.float32)
np.argmax(testArr, axis=1)
>>> array([1, 3, 0, 1, 1])