python numpyの配列内の最大値のインデックス番号を取得する

python numpyの配列内の最大値のインデックス番号を取得する

pythonで、ライブラリnumpyのargmaxを使用して、配列内の最大値のインデックス番号を取得するサンプルコードを記述してます。pythonのバージョンは3.8.5を使用してます。

環境

  • OS windows10 pro 64bit
  • python 3.8.5

numpyインストール

numpyをインストールされていない方は、pipでインストールしておきます。

pip install numpy

# Successfully installed numpy-1.19.4

argmax使い方

argmaxを使用すると、配列内の最大値のインデックス番号を取得することが可能です。

import numpy as np

arr = np.array([2, 3, 1, 4, 5, 6])

print(np.argmax(arr))
# 5

# 同じ値の場合は、最初の方のインデックス番号を返す
arr = np.array([2, 6, 1, 4, 5, 6])

print(np.argmax(arr))
# 1

2次元配列の場合は、フラットになった状態の値が返ってきます。

import numpy as np

arr = np.array([[1, 2, 3],[4, 5, 6],[7, 8, 9]])

# フラット化された値が返ってきます
print(np.argmax(arr))
# 8

# unravel_indexを使用すると8ではなく(2,2)で取得することが可能です
print(np.unravel_index(np.argmax(arr, axis=None), arr.shape))
# (2, 2)