動的に条件を作ってpandas.DataFrameから値を取得したいときがある。次の例を考えてみる。
import pandas as pd df = pd.DataFrame({ "A" : [1, 0, 1, 0], "B" : ["a", "a", "b", "b"], "C" : [1, 2, 3, 4] })
このとき、ある時は 「A == 0」、また別の時は 「A == 0 かつ B == 0」の条件でDataFrameの行を取得したいときはないだろうか?私はある。
こんなとき、リストの要素を結合するjoinとDataFrameから条件抽出するための関数queryを使うと解決できる。
まず、joinを使うと次のような操作をリストに対して実行できる。
test = ["a", "b", "c"] "&".join(test) ''' 出力結果 'a&b&c' '''
すなわち、joinの前につけた文字列でリストの要素を結合することができる。
また、queryは次のような働きをする。
df.query("A==0 & B==0")
このときの出力結果は、図1のようになる。つまり、queryの中にdfのカラム名と各カラムに対する条件を付与することで、条件に合致する行を取得できる。
これらを組み合わせて、カラムを指定したら条件の数に関わらず行を抽出する関数を作成する。ただし、今回は簡単のため、各条件は「カラム==0」であり、これらのアンド条件を作成する関数とする。他の条件を作成したい方は、適当に調整してください。
def get_row(target_df, cond_columns_list): ''' df : pd.DataFrame, 抽出対象のDataFrame cond_list : 文字列のlist, 条件を付与したいカラム ''' # query作成 cond_list = ["({c} == 0)".format(c=column) for column in cond_columns_list] q = "&".join(cond_list) # 条件抽出 tmp_df = target_df.query(q) return tmp_df
実際に使ってみるとこんな感じ。
get_row(target_df=df, cond_columns_list=["A"]) # →図2
get_row(target_df=df, cond_columns_list=["A", "B"]) # →図3
無事に取得できた!