NumPyのargmax関数は、多次元配列の中の最大値の要素を持つインデックスを返す関数です。np.maxを使うと、最大値の要素を返すことができました。argmaxは、最大の要素のインデックスを返します。maxの使い方については、以下の記事を参考にしてください。

最大値を抜き出すnumpy.amaxとndarray.max関数の使い方 /features/numpy-max.html

本記事では、最大値のインデックスを返すnp.argmaxとnp.ndarray.argmaxの２つの関数について解説します。

np.ndarray.argmax 、np.argmax

まずはnp.ndarray.argmaxとnp.argmax両方のAPIドキュメントを見ていきましょう。

np.ndarray.argmax(axis = None, out = None)

params:

パラメータ名 型 概要
axis int (省略可能)初期値None
最大値を読み取る軸の方向を指定します。
out array_like
(配列に相当するもの)		 (省略可能)初期値None
返されたインデックスの値を収納する配列を指定します。

returns:

指定された配列の中で最大値となっている要素のうち先頭のインデックスを返します。
軸が指定されている場合は、その軸に沿った向きにおける範囲でのインデックスを返します。

np.argmax(a, axis = None, out = None)

params:

パラメータ名 型 概要
a array_like
(配列に相当するもの)		 最大値のインデックスを知りたい配列をここで指定します。
axis int (省略可能)初期値None
最大値を読み取る軸の方向を指定します。
out array_like
(配列に相当するもの)		 (省略可能)初期値None
返されたインデックスの値を収納する配列を指定します。

returns:

指定された配列の中で最大値となっている要素のうち先頭のインデックスを返します。
軸が指定されている場合は、その軸に沿った向きにおける範囲でのインデックスを返します。

使い方はnp.argmaxの方は、第一引数に最大値を取得したい配列を指定します。一方で、ndarray.argmaxの方はメソッドのようにして呼び出します。

最後の引数のoutはあまり使用しません。出力の配列を予め作っておいた配列にしたい場合には、指定します。

まずは、1次元配列の例で動作を確認してみます。

In [1]: import numpy as np

In [2]: a = np.random.randint(10, size = 10) # まずは1次元配列から。

In [3]: a  # aの値を確認する。
Out[3]: array([2, 3, 3, 1, 4, 4, 5, 0, 5, 4])

In [4]: np.argmax(a)
Out[4]: 6

In [5]: a.argmax()
Out[5]: 6

この例では、ランダムで生成した配列の中身が[2, 3, 3, 1, 4, 4, 5, 0, 5, 4]となっています。最大値は5ですが、インデックスは6と8の二回出現します。こういった場合には、最初の最大値のインデックスを返します。

次は多次元配列にして動作を確認します。axisを指定しない場合には、1次元のときと同様に最大値の要素のインデックスを1次元配列に直した時のインデックスで返します。

In [6]: b = np.random.randint(10, size = (3,4)) # 今度は3×4の2次元配列を生成。

In [7]: b # bの要素を確認。
Out[7]:
array([[4, 9, 1, 5],
       [3, 5, 6, 2],
       [9, 8, 1, 0]])

In [8]: np.argmax(b) # 2次元配列の最大値のインデックスを求めているが、1次元配列に直したあとのインデックスが返ってくる。この場合は1。
Out[8]: 1

In [9]: b.argmax() # np.ndarray.argmax() についても同様。
Out[9]: 1

今回生成した多次元配列の最大値は9ですね。9の最初のインデックスは1になります。

次は軸(axis)を指定します。

In [10]: b
Out[10]:
array([[4, 9, 1, 5],
       [3, 5, 6, 2],
       [9, 8, 1, 0]])

In [11]: np.argmax(b, axis = 0) # axis = 0(この場合は行)方向の最大値を見ていく。(縦方向に見てその最大値のインデックスをとっていくから要素数は4)
Out[11]: array([2, 0, 1, 0])

In [12]: b.argmax(axis = 0) # np.ndarray.argmaxも同様。
Out[12]: array([2, 0, 1, 0])

今回の配列の形状は(3, 4)なので、axis = 0を指定した場合には、行方向に最大のインデックスを返すことになります。行方向で最大の要素を探すと以下のようになります。 行方向ということなので、列ごとの最大値のインデックスを返します。

b[?, 0] → 9 (b[2, 0])
b[?, 1] → 9 (b[0, 1])
b[?, 2] → 6 (b[1, 2])
b[?, 3] → 5 (b[0, 3])

上記の?の部分がインデックスになります。つまりarray([2, 0, 1, 0])が出力されることになります。同様にaxis = 1を指定すると、行方向に最大の要素のインデックスを返します。

In [13]: np.argmax(b, axis = 1) # axisを1に指定して見ます。このときは列方向の最大値を見ていきます。(横方向の成分の中での最大値)
Out[13]: array([1, 2, 0])

In [14]: b.argmax(axis = 1)
Out[14]: array([1, 2, 0])

3次元配列における軸の指定方法も確認してみます。

In [15]: c = np.random.randint(10, size = (2,3,4)) # 2×3×4の3次元行列を生成。

In [16]: c # cの値を確認。
Out[16]:
array([[[7, 8, 9, 9],
        [5, 3, 8, 6],
        [4, 9, 7, 3]],

       [[1, 3, 4, 3],
        [7, 7, 9, 0],
        [2, 9, 5, 6]]])

In [17]: np.argmax(c, axis = 0)
Out[17]:
array([[0, 0, 0, 0],
       [1, 1, 1, 0],
       [0, 0, 0, 1]])

この例では、axis= 0を指定した場合には次のようになります。

c[?, 0, 0] → 7 (c[0, 0, 0])
c[?, 0, 1] → 8 (c[0, 0, 1])
c[?, 0, 2] → 9 (c[0, 0, 2])
c[?, 0, 3] → 9 (c[0, 0, 3])

c[?, 1, 0] → 7 (c[1, 1, 0])
c[?, 1, 1] → 7 (c[1, 1, 1])
c[?, 1, 2] → 9 (c[1, 1, 2])
c[?, 1, 3] → 6 (c[0, 1, 3])

c[?, 2, 0] → 4 (c[0, 2, 0])
c[?, 2, 1] → 9 (c[0, 2, 1])
c[?, 2, 2] → 7 (c[0, 2, 2])
c[?, 2, 3] → 6 (c[1, 2, 3])

他の軸も指定してみます。

In [18]: c.argmax(axis = 0)
Out[18]:
array([[0, 0, 0, 0],
       [1, 1, 1, 0],
       [0, 0, 0, 1]])

In [19]: np.argmax(c, axis = 1)
Out[19]:
array([[0, 2, 0, 0],
       [1, 2, 1, 2]])

In [20]: c.argmax(axis = 1)
Out[20]:
array([[0, 2, 0, 0],
       [1, 2, 1, 2]])

In [21]: np.argmax(c, axis = 2)
Out[21]:
array([[2, 2, 1],
       [2, 2, 1]])

In [22]: c.argmax(axis = 2)
Out[22]:
array([[2, 2, 1],
       [2, 2, 1]])

また、最小値のインデックスを取得したい場合はargminを使うことで取得することができます。