Pandasで上位n番目までの値を得る

やること

Pandasで上位n番目までの値を得る | 自調自考の旅という記事を発見しました。

PandasのDataFrameを使ってデータを処理すると、maxやminといったメソッドを用いて簡単に各カラムの最大値や最小値を得ることが出来ます。しかしながら、現時点（pandas ver 1.1.2）では、2番目の最大値や最小値、3番目の最大値や最小値、、、といった値を得る機能は備わっていません。［……］そこで本記事では、なるべく少ない行数でかけて、かつ以下図のイメージのようにDataFrameの各カラムの上位n番目までの値を得られるようなスクリプトを紹介します。

つまり、以下のようなデータフレームdfがあったときに、

import numpy as np
import pandas as pd

np.random.seed(0)
df = pd.DataFrame(np.random.permutation(50).reshape(10, 5))
df

	0	1	2	3	4
0	28	11	10	41	2
1	27	38	31	22	4
2	33	35	26	34	18
3	7	14	45	48	29
4	15	30	32	16	42
5	20	43	8	13	25
6	5	17	40	49	1
7	12	37	24	6	23
8	36	21	19	9	39
9	46	3	0	47	44

各列に対して上位（あるいは下位）３つを抜き出して、以下のようなデータフレームを取得するような関数・メソッドがpandasにはないということです。

	0	1	2	3	4
1	46	43	45	49	44
2	36	38	40	48	42
3	33	37	32	47	39

方針

いま例に挙げているデータフレームdfはすべての列が同じデータ型になっていますが、実際のデータは列毎にデータ型が違うことがおおいので、df.apply()で列毎に処理することを考えます。

つまり、以下のようなシリーズsを渡して、

	0
0	28
1	27
2	33
3	7
4	15
5	20
6	5
7	12
8	36
9	46

以下のようなシリーズが返却される関数を作っていきます。

	0
1	46
2	36
3	33

方法

引用した記事では、以下のような関数が示されています。なお、オプション引数は以下のとおりです。

topnum: 取得する件数。デフォルトは3。
getmin: Trueにすると昇順で取得。デフォルトは降順。
getindex: Trueにすると値ではなくインデックスを返す。

def getmax(series, topnum=3, getmin=False, getindex=False):
    if getindex is False:
        series = (series.sort_values(ascending=getmin).head(topnum)
                  .reset_index(drop=True))
        series.index += 1
        return series
    else:
        return series.sort_values(ascending=getmin).head(topnum).index

この方法は、まずシリーズ全体をソートしたのちに、上側を取得する、そしてインデックスを1始まりにするために、インデックスをリセットした後に1を加えるという作業を行っています。

しかし、これはないでしょう。

pandasでシリーズから上位n件を取得するには、pd.Series.nlargset()メソッド（およびpd.Series.nsmallest()メソッド）が最適解です。

%timeit df[0].sort_values(ascending=False).head(3)
%timeit df[0].nlargest(3)

	0
9	46
8	36
2	33

299 µs ± 9.43 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
153 µs ± 1.58 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

速度も倍違います（なお、速度を考えるとNumPyを使うことも考えに入れられますが、コードの簡潔さがやや失われます）。

1始まりの連番インデックスをつけるのについてもかなり冗長です。直接1始まりの連番を設定すればいいところ、0始まりのインデックスにしてから1を足しています。np.arange()やrange()を使うことを思いつくでしょうが、pd.RangeIndex()を使いましょう。

test_s = pd._testing.makeStringSeries(10000)

%timeit s2 = test_s.reset_index(drop=True); s2.index += 1
%timeit s2 = test_s.set_axis(range(1, len(test_s)+1))
%timeit s2 = test_s.set_axis(np.arange(1, len(test_s)+1))
%timeit s2 = test_s.set_axis(pd.RangeIndex(1, len(test_s)+1))

109 µs ± 2.98 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
40.2 µs ± 1.02 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
64.7 µs ± 1.49 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
39.8 µs ± 931 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

結論

そういうわけで以下のようになります。

def getmax_rev(series, topnum=3, getmin=False, getindex=False):
    out = series.nsmallest(topnum) if getmin else series.nlargest(topnum)
    return out.index if getindex else out.set_axis(pd.RangeIndex(1, topnum+1))