Python numpy argmax用法及代碼示例

本文簡要介紹 python 語言中 numpy.argmax 的用法。

用法:

numpy.argmax(a, axis=None, out=None, *, keepdims=<no value>)

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

參數：

a： array_like

輸入數組。

axis： 整數，可選

默認情況下，索引位於展平數組中，否則沿指定軸。

out： 數組，可選

如果提供，結果將被插入到這個數組中。它應該具有適當的形狀和數據類型。

keepdims： 布爾型，可選

如果將其設置為 True，則縮小的軸將作為尺寸為 1 的尺寸留在結果中。使用此選項，結果將針對數組正確廣播。

返回：

index_array： 整數數組

數組中的索引數組。它與 a.shape 具有相同的形狀，但刪除了沿軸的尺寸。如果 keepdims 設置為 True，則 axis 的大小將為 1，結果數組的形狀與 a.shape 相同。

注意：

如果最大值多次出現，則返回與第一次出現對應的索引。

例子：

>>> a = np.arange(6).reshape(2,3) + 10
>>> a
array([[10, 11, 12],
       [13, 14, 15]])
>>> np.argmax(a)
5
>>> np.argmax(a, axis=0)
array([1, 1, 1])
>>> np.argmax(a, axis=1)
array([2, 2])

N 維數組的最大元素的索引：

>>> ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
>>> ind
(1, 2)
>>> a[ind]
15

>>> b = np.arange(6)
>>> b[1] = 5
>>> b
array([0, 5, 2, 3, 4, 5])
>>> np.argmax(b)  # Only the first occurrence is returned.
1

>>> x = np.array([[4,2,3], [1,0,3]])
>>> index_array = np.argmax(x, axis=-1)
>>> # Same as np.amax(x, axis=-1, keepdims=True)
>>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1)
array([[4],
       [3]])
>>> # Same as np.amax(x, axis=-1)
>>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1)
array([4, 3])

將 keepdims 設置為 True，

>>> x = np.arange(24).reshape((2, 3, 4))
>>> res = np.argmax(x, axis=1, keepdims=True)
>>> res.shape
(2, 1, 4)