pandasのSeriesのlocには関数も渡せる

何となくpandasのドキュメントを眺めていたら見つけた小ネタの紹介です。
この記事を読むと、pandasのSeriesをもっと手軽に値で絞り込める様になります。
参考: pandas.Series.loc — pandas 2.0.3 documentation

pandasのlocといえば、自分としてはDataFrameで使うことが多いプロパティですが、もちろんSeriesにも実装されています。そして、これを使うとindexの値に従って要素を絞り込むことができます。

今回見つけたのは、このlocにcallable、要するにメソッドが渡せるってことです。渡したメソッドにSeriesの値が渡され、その結果がTrueのものに絞り込まれます。

こんな感じで使えます。例えば値が3以上の要素だけに絞り込む例です。

import pandas as pd

# Seriesを作成
s = pd.Series([1, 2, 3, 4, 5], index=['a', 'b', 'c', 'd', 'e'])

# 値が3以上の要素だけを絞り込む
s[lambda x: x>=3]
"""
c    3
d    4
e    5
dtype: int64
"""

上記の例だとメリットがわかりにくいかと思います。と言うのもpandasに慣れてる人だったら次の様に書けることがわかってると思いますし、タイプ数も少なく可読性も高いからです。

s[s>=3]
"""
c    3
d    4
e    5
dtype: int64
"""

単純な大小比較や一致不一致ではなく、もっと複雑な判定を行うメソッドを適用するときなどは、このlocにメソッドを渡すやり方が便利に使えますね。

僕が個人的に気に入ったのは、この絞り込みが特定の変数に格納されていないSeriesについても使えると言うことです。
「特定の変数に格納されていないSeries」ってのは、例えばDataFrameのvalues_counts()メソッドなどを実行した結果として得られる値などの形で取得されます。

例えば、dfというDataFrameがあるとして、そのcolumn_name列の値の出現回数を数え、そのうち出現回数が10以上のものだけを取り出すとしましょう。

通常であれば、value_counts()の結果を何かの変数に格納して実行するか、もしくは2回value_counts()を実行する非効率に我慢するかして次の様に実装します。

# カウント結果を一度変数に格納する場合
count_sr = df["column_name"].value_counts()
count_sr[count_sr >= 10]

# value_countsを2回実行する場合
df["column_name"].value_counts()[df["column_name"].value_counts() >= 10]

これが、locにメソッドが渡せることを知っていると次の様に書けます。

# 無駄な変数も定義しないし、value_counts()の実行も1回でいい書き方
df["column_name"].value_counts()[lambda x: x>=10]

個人的に、「ある列の要素ごとに数を数えて、一定件数以上データがあったものだけ残す処理」ってのをやることが頻繁にあり、Seriesのfiterメソッドが値ではなくindexにしか使えないことを日々残念に思っていた自分にとってはめっちゃ嬉しいテクニックだったので紹介しました。

ちなみに、DataFrameのlocも同じ様にcallableを渡せます。こっちはあまり使い道が思いつかないですね。
参考: pandas.DataFrame.loc — pandas 2.0.3 documentation

コメントを残す

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