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

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

## やること

[Pandasで上位n番目までの値を得る | 自調自考の旅](https://own-search-and-study.xyz/2020/10/24/pandas-dataframe-second-max/)という記事を発見しました。

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

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

```python
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`にすると値ではなくインデックスを返す。

```python
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
```

この方法は、まずシリーズ全体を[ソート](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.sort_values.html)したのちに、[上側を取得](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.head.html)する、そしてインデックスを1始まりにするために、[インデックスをリセット](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.reset_index.html)した後に1を加えるという作業を行っています。

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

pandasでシリーズから上位n件を取得するには、<code>[pd.Series.nlargset()](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.nlargest.html)</code>メソッド（および<code>[pd.Series.nsmallest()](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.nsmallest.html)</code>メソッド）が最適解です。

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

|       |    0 |
| ----: | ---: |
| **9** |   46 |
| **8** |   36 |
| **2** |   33 |

```text
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()`を使うことを思いつくでしょうが、<code>[pd.RangeIndex()](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.RangeIndex.html)</code>を使いましょう。

```python
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))
```

```text
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)
```

## 結論

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

```python
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))
```