numpyで配列生成などはよくやりますね。
その配列の最大値のインデックスを知りたい、なんてこともよくありますよね。ないなんて言わせません
argmax関数
numpyには最大値のインデックスを求めてくれるargmax関数がありますが、これには少し問題があって最大値が複数あると最も小さいインデックスのみを返してくれます。
このことについては公式サイトのマニュアルにも末尾に記載されています。
test.py
import numpy as np
b = np.arange(6) # [0, 1, 2, 3, 4, 5]
b[1] = 5 # [0, 5, 2, 3, 4, 5]
np.argmax(b) # 1
for文&max関数
対策としてはfor文でやる方法もありますが、スマートじゃありませんし時間がかかりますよね。
findMax.py
a = np.array([0, 5, 2, 3, 5])
maxValue = np.max(a)
maxIndex = []
for i in range(len(a)):
if maxVaule == a[i]:
maxIndex.append(i)
print(maxIndex) # [1, 4]
enumerate関数&内包表記
そこでenumerate関数と内包表記を用いてカッコよく最大値を複数取り出します。
下記にそれぞれの参考サイトを記載します。
Python zipとenumerateの使い方
pythonの内包表記を少し詳しく
では早速書いてみましょう。
findMax2.py
a = np.array([0, 5, 2, 3, 5])
maxIndex = [i for i, x in enumerate(a) if x == max(a)]
print(maxIndex) # [1, 4]
これだけです。とても簡潔でいいですよね。
これは余談ですが内包表記を使ってリストのstr型をリストのint型に１行で変換することもできちゃいます。paizaなどのプログラムスキルチェックしてくれるサイトの問題等で使えますね。
strToInt.py
strList = ["0", "1", "2", "3"]
intList = [int(i) for i in strList]
print(intList) # [0, 1, 2, 3]