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

Question

I have a dataframe like this for example:

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

What I need is to filter the df based on the value in column 'A'. The problem is that the values to filter by are supplied by the end user. For example:

cond = ['a', 'b']

means that the user wants to filter the df and keep all values 'a' and 'b' in the column 'A'. So in this case I'll need to filter the df with this condition:

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

But the next time the values in the cond list can be different and I need to account for it. So far I've tried the for loop. I was pretty certain it wasn't going to work... and it didn't:

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

I've also tried to create a generator under df.query() and had high hopes for this, but it didn't work either. Unfortunately, the method doesn't accept generators:

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

Not quite sure what else to try. Has anyone dealt with this type of problem before?

Answer 1

你可以试试

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

Answer 2

也可以尝试替代@BEN_YO

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