Python
pandas

pandasのDataFrameでargmaxを取りたい時

pandasのDataFrameを使って、argmaxのようなしたいとき、どのようにすれば良いか。

これを1発でやる関数がpandasには備わっています。それは、idxmax関数です。

例えば、以下のようなデータで、各列について、最大値をとるindex番号が欲しいとする。

0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
0 0.189460 0.086207 0.203066 0.208812 0.259933 0.185123 0.158537 0.200099 0.149425 0.138298 0.185732 0.203560 0.191906 0.199288 0.205338 0.203721 0.203377 0.158458 0.293388 0.244068
1 0.209651 0.256466 0.238045 0.132184 0.134794 0.230882 0.208426 0.211899 0.272578 0.328267 0.138991 0.171301 0.216710 0.192171 0.217290 0.176644 0.199074 0.235546 0.107438 0.144068
2 0.196588 0.153017 0.171099 0.270115 0.205347 0.157334 0.216741 0.185818 0.182266 0.177812 0.242312 0.299221 0.147520 0.277580 0.152408 0.216799 0.199033 0.244111 0.260331 0.223729
3 0.215362 0.321121 0.190227 0.195402 0.157817 0.177183 0.161308 0.219288 0.177340 0.200608 0.234932 0.157953 0.255875 0.256228 0.235622 0.185117 0.198436 0.231263 0.185950 0.210169
4 0.188939 0.183190 0.197563 0.193487 0.242109 0.249478 0.254989 0.182896 0.218391 0.155015 0.198032 0.167964 0.187990 0.074733 0.189342 0.217720 0.200080 0.130621 0.152893 0.177966

これは以下のようにすれば、一発!

>> df = pd.DataFrame('hogehoge')
>> df.idxmax()

すると、以下のような結果になります。

0     3
1     3
2     1
3     2
4     0
5     4
6     4
7     3
8     1
9     1
10    2
11    2
12    3
13    2
14    3
15    4
16    0
17    2
18    0
19    0
dtype: int64

ちなみに返り値のデータ型は、Seiresです。

>> type(df.idxmax())
<class 'pandas.core.series.Series'>

参考

https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.idxmax.html