在 Pandas 中使用条件列表过滤 DataFrame

Posted 2023-03-14

【中文标题】在 Pandas 中使用条件列表过滤 DataFrame

【英文标题】：Using a list of conditions to filter a DataFrame in Pandas

【发布时间】：2017-08-26 21:49:06

【问题描述】：

我希望有一个函数，它接受一个任意长度的条件列表，并在所有条件之间放置一个 & 符号。示例代码如下。

df = pd.DataFrame(columns=['Sample', 'DP','GQ', 'AB'],
         data=[
               ['HG_12_34', 200, 35, 0.4],
               ['HG_12_34_2', 50, 45, 0.9],
               ['KD_89_9', 76, 67, 0.7],
               ['KD_98_9_2', 4, 78, 0.02],
               ['LG_3_45', 90, 3, 0.8],
               ['LG_3_45_2', 15, 12, 0.9]
               ])


def some_func(df, cond_list):

    # wrap ampersand between multiple conditions
    all_conds = ?

    return df[all_conds]

cond1 = df['DP'] > 40
cond2 = df['GQ'] > 40
cond3 = df['AB'] < 0.4


some_func(df, [cond1, cond2]) # should return df[cond1 & cond2]
some_func(df, [cond1, cond3, cond2]) # should return df[cond1 & cond3 & cond2]

我将不胜感激。

【参考方案1】：

您可以为此使用functools.reduce：

from functools import reduce

def some_func(df, cond_list):
    return df[reduce(lambda x,y: x&y, cond_list)]

或者，就像@AryaMcCarthy 所说，您可以使用运算符包中的and_：

from functools import reduce
from operator import and_

def some_func(df, cond_list):
    return df[reduce(and_, cond_list)]

或者用 numpy - 就像@ayhan 说的 - 它也有一个逻辑和减少：

from numpy import logical_and

def some_func(df, cond_list):
    return df[logical_and.reduce(cond_list)]

所有三个版本都会为您的示例输入产生以下输出：

>>> some_func(df, [cond1, cond2])
       Sample  DP  GQ   AB
1  HG_12_34_2  50  45  0.9
2     KD_89_9  76  67  0.7
>>> some_func(df, [cond1, cond2, cond3])
Empty DataFrame
Columns: [Sample, DP, GQ, AB]
Index: []

【讨论】：

使用operator.and_ 代替您的自定义 lambda 可能会更好。

@AryaMcCarthy：是的，确实更整洁。

或者，来自 numpy：np.logical_and.reduce([cond1, cond2, cond3])