根据可变数量的条件过滤数据框

时间:2020-08-26 01:24:42

标签: python pandas

例如,我有一个这样的数据框:

df = pd.DataFrame({'A':['a', 'a', 'b', 'c', 'a', 'b',], 'B': [1, 2, 3, 4, 5, 6,]})

我需要根据'A'列中的值过滤df。问题是最终用户要提供要过滤的值。例如:

cond = ['a', 'b']

表示用户希望过滤df,并将所有值“ a”和“ b”保留在“ A”列中。因此,在这种情况下,我需要使用以下条件过滤df:

df = df.loc[(df['A'] == 'a') | (df['A'] == 'b')]

但是下一次cond列表中的值可能会有所不同,我需要考虑一下。到目前为止,我已经尝试过for循环。我敢肯定这是行不通的...而且没有:

for item in cond:
    df = df.loc[df['A'] == item]

我也曾尝试在df.query()下创建一个生成器,对此寄予了厚望,但它也不起作用。不幸的是,该方法不接受生成器:

df = df.query(f'A == {x}' for x in cond)
# or
df = df.query('A == @x' for x in cond)

不太确定还有什么尝试。以前有没有人处理过这类问题?

2 个答案:

答案 0 :(得分:2)

您可以尝试

df = df.loc[df['A'].isin(cond)]

答案 1 :(得分:1)

也可以尝试替代@BEN_YO

 df.query('A==@cond')