pandasのメソッドで、上位n件や下位n件のデータを取得する

先日紹介したbar chart raceのライブラリのドキュメントやソースコードを読んでいて、その中で nlargest というメソッドを見つけたのでその紹介です。その対となる nsmallest というメソッドもあります。

これが何をするメソッドとかというと、DataFrameやSeriesのデータの値が大きい方からn件(nlargest)や小さい方からn件(nsmallest)を取得してくれるものです。
え、sort_values() して、 head(n)やtail(n)すればいいじゃん、という声も聞こえてきそうですし、実際僕もそう思ってるのですが、多少の利点がちゃんとあるので読んでいただければ幸いです。

公式ドキュメントはこちらになります。
pandas.DataFrame.nlargest
pandas.DataFrame.nsmallest
pandas.Series.nlargest
pandas.Series.nsmallest

使い方は簡単で、Seriesの方であれば、取得したい件数を最初の引数nに渡してあげるだけ、DataFrameの方は、取得したい件数と合わせて、どの列の上位/下位を取得したのかを2つ目の引数columnsに渡してあげればOKです。

とりあえず、適当に作ったDataFrameに対して適当に列を指定して5項目ほど取得してみましょう。

import pandas as pd
import numpy as np


# 50行3列の乱数データを生成する
data = np.random.randint(1, 50, size=(50, 3))
df = pd.DataFrame(data, columns=["col1", "col2", "col3"])
print(df.shape)
# (50, 3)

print(df.nlargest(5, "col2"))
"""
    col1  col2  col3
46    30    48    28
17    47    47    31
33     9    45    30
16     3    44    33
26    16    44     2
"""

見ての通り、指定した”col2″でソートした上でその値が大きいものから順番に、5項目選択されています。

nlargest/ nsmallest にはもう一つ、keepという引数があります。これは、値が等し鋳物が複数あって、n位にランクインするものが一意に決められないときにその取り扱いを指定するものです。
“first”(デフォルト)を指定すると、元のデータで先に登場指定したものが優先され、”last”を指定すると、最後に登場したものが優先されます。また、”all”にすると、同率だったものが全部含まれます。

print(df.nlargest(5, "col2", keep="first"))
"""
    col1  col2  col3
46    30    48    28
17    47    47    31
33     9    45    30
16     3    44    33
26    16    44     2
"""

print(df.nlargest(5, "col2", keep="last"))
"""
    col1  col2  col3
46    30    48    28
17    47    47    31
33     9    45    30
40    48    44     1
26    16    44     2
"""

print(df.nlargest(5, "col2", keep="all"))
"""
    col1  col2  col3
46    30    48    28
17    47    47    31
33     9    45    30
16     3    44    33
26    16    44     2
40    48    44     1
"""

“col2″に値が44のレコードが3つ存在するのですが、”first”と”last”で選択されたレコードが違うのがわかりますね。そして”all”を指定すると3レコードとも返され、結果が6行になっています。

このkeep引数が存在することのほか、sort_values/head に比べると、速度面でも優れているそうです。

This method is equivalent to df.sort_values(columns, ascending=False).head(n), but more performant.

とドキュメントにもあります。
ソースを読んで無いので予想ですが、sort_values/headの方は最終的な結果に必要ない行まで全部ソートを完了させるに対して、nlargest/nsmallestの方は必要なデータだけ並べ替えてソートを打ち切ってるのではないかと思っています。

コードの実行例は載せませんでしたが、nsmallestもnlargestと同じように使うことができ、こちらは結果が小さい順に取得されます。

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です