NumPyの最小値のインデックスを取得するargmin()の使い方

NumPyのargmin関数は、配列内の最小値のインデックスを返す関数です。全く同じ機能のものに、ndarray.argmin()メソッドも用意されています。

通常、NumPyの配列から最小値を取得するにはnp.amin()を使います。しかし、np.argmin()で取得できるインダイスを使うと、より複雑な操作を行うことが可能です。NumPyの配列には、Python標準のリストやタプルと比べて、豊富なスライスのテクニックが用意されているからです。そのため、np.argmin()を使いこなすには、インダイスの配列を使ったスライスのテクニックを身につけておく必要があります。

それでは、このページでは、np.argmin()の使い方についてサンプルコードを見ながら確認していきます。

1. 書式

まずは書き方と引数を確認しましょう。

numpy.argmin(a, axis=None, out=None)

ndarray.argmin(axis=None, out=None)

2. サンプルコード

実際のコードを見ながら使い方を確認していきましょう。

なお、np.argmin関数とndarray.argminメソッドは機能は全く同じなので、ここではnp.argmin関数で解説していきます。

1次元配列の場合

まずは、以下の1次元配列を例に見てみましょう。

import numpy as np
rng = np.random.default_rng()
arr = rng.integers(0, 10, (10, ))
arr

array([9, 2, 7, 9, 9, 1, 6, 5, 0, 0])

この配列を関数に渡すと、最小値のインデックスを返します。

# 最小値のインデックスを取得
np.argmin(arr)

8

スライスすると値を参照することができます。スライスについては『NumPyの配列のスライスの必須テクニックまとめ』で解説しています。

# スライスして値を取得
arr[8]

0

確かに、最小値である0を参照しています。なお、この例のように最小値が複数ある場合は最初のインデックスを返します。

多次元配列の場合

多次元配列を渡した時は、その多次元配列を1次元配列化(flatten)した時のインデックスを返します。

以下のコードでご確認ください。

import numpy as np
rng = np.random.default_rng()
arr = rng.integers(0, 10, (2, 3))
arr

array([[2, 3, 9],
       [3, 1, 2]])

# 最小値のインデックスを取得
np.argmin(arr)

4

この場合、値を参照するには、元の多次元配列をndarray.flattenメソッドで1次元配列化してからスライスします。

# スライスして値を取得
arr.flatten()[4]

1

または、以下のように、インダイスを任意のshapeに対応する形に変換してくれるunravel_index()を使う方法もあります。

#  元の2次元配列arr のshapeに対応するインダイスを取得
ind = np.unravel_index(np.argmin(arr, axis=None), arr.shape)
ind

(1, 1)

arr[ind]

1

unravel_index()は、第一引数にインダイスの配列を渡し、第二引数で変換したいshapeを指定します。インダイスを扱う上で非常に便利なので、一緒に覚えておくと良いでしょう。

次元軸を指定

オプション引数axis=を使うと、要素を検索する次元軸を指定することができます。以下の2次元配列を使って確認しましょう。

import numpy as np
rng = np.random.default_rng()
arr = rng.integers(0, 10, (2, 3))
arr

array([[7, 9, 0],
       [5, 1, 5]])

axis=0は最初の軸を指定することと同じです。2次元配列では、最初の軸は2次元軸（縦軸）に該当します。

以下で確認しましょう。

# 2次元軸(縦軸)の最小値のインデックスを取得
ind = np.argmin(arr, axis=0)
ind

array([1, 1, 0])

このインダイスから各列の最小値を一括取得するには、以下のようにtake_along_axis()を使います。

# 2次元軸（縦軸)ごとの最小値を一括取得
ind_array = np.expand_dims(ind, axis=0) #インデックス配列の次元数をaに合わせる
np.take_along_axis(arr, ind_array, axis=0)

array([[5, 1, 0]])

expand_dimsは配列の次元数を増やす関数です。これによって、元の配列a の次元数とインダイス配列ind の次元数を揃えてブロードキャスト可能にしています。np.take_arlong_axisは多次元配列から指定のインデックス番号の要素を一括して取得する関数です。こちらの関数では、元の配列a とインダイス配列の次元数が揃っている必要があります。

axis=1を指定した時は、1次元軸（横軸）の要素から最小値のインデックスを取得します。

# 1次元軸(横軸)の最小値のインデックスを取得
ind = np.argmin(arr, axis=-1)
ind

array([2, 1])

1次元軸ごとの最小値を一括取得するには、同じくtake_along_axis()を使います。

# 1次元軸(横軸)ごとの最小値を一括取得
ind_array = np.expand_dims(ind, axis=-1) #インデックス配列の次元数をaに合わせる
np.take_along_axis(arr, ind_array, axis=-1)

array([[0],
       [1]])

3次元配列以上になっても考え方は同じです。

戻り値で既存の配列を上書き

使うことはほぼありませんが、オプション引数 out= で、既存の配列を戻り値の配列で上書きすることができます。この時、お互いの配列とデータ型の形状が一致している必要があります。

以下のコードでご確認ください。

# 配列を作成
import numpy as np
rng = np.random.default_rng()
arr = rng.integers(5, size=(2, 5))
arr

array([[2, 1, 4, 3, 3],
       [4, 2, 2, 0, 2]])

# 上書きする配列を作成
x = np.zeros(5, dtype=int)
x

array([0, 0, 0, 0, 0])

# np.argminの戻り値でxを上書き
np.argmax(arr, axis=0, out=x)
x

array([0, 0, 1, 1, 1])

3. まとめ

以上が、np.argmin関数の使い方です。

ページ中で解説してた np.expand_dims や np.take_along_axis、np.unravel_index の書き方もしっかり頭に入れておきましょう。

なお、ndarray.argminメソッドも使い方はまったく同じです。念のためサンプルコードを貼っておきます。

import numpy as np
#  配列を生成
rng = np.random.default_rng()
arr = rng.integers(0, 10, (2, 3))
print('arr:\n', arr, '\n')

#  メソッドなので配列を前に書く
print('arr.argmin()\n', arr.argmin(), '\n')

#  axisの指定
print('arr.argmin(axis=0)\n', arr.argmin(axis=0), '\n')

#  outの指定
a_out = np.zeros(2, dtype=int)
arr.argmin(axis=1, out=a_out)
print('a_out:\n', a_out)

arr:
 [[7 1 6]
 [1 1 2]] 

arr.argmin()
 1 

arr.argmin(axis=0)
 [1 0 1] 

a_out:
 [1 0]