初心者データサイエンティストの備忘録

調べたことは全部ここに書いて自分の辞書を作る

動的に条件を作ってpandas.DataFrameから値を取得する

 動的に条件を作って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のカラム名と各カラムに対する条件を付与することで、条件に合致する行を取得できる。

f:id:aisinkakura_datascientist:20220311020005j:plain:w300
図1:抽出結果

 これらを組み合わせて、カラムを指定したら条件の数に関わらず行を抽出する関数を作成する。ただし、今回は簡単のため、各条件は「カラム==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

f:id:aisinkakura_datascientist:20220311021309j:plain:w300
図2: 「A==0」で抽出した結果

get_row(target_df=df, cond_columns_list=["A", "B"])

# →図3

f:id:aisinkakura_datascientist:20220311020005j:plain:w30
図3:「A==0 & B==0」で抽出した結果

 無事に取得できた!