python - 在 pandas Dataframe 中索引和选择数据时动态更新条件

Question

从数据框中选择数据时，我需要按照输入列表的长度改变条件的方式。这是我当前的代码。列表中的第一个元素是列的名称，第二个元素是列的值。

import pandas as pd

list_1 = [('a', 2), ('b', 5)]
list_2 = [('a', 1), ('b', 2), ('c', 3)]

data = pd.DataFrame([[1, 2, 3], [1, 2, 3], [1, 1, 1], [2, 5, 6]], columns=['a', 'b', 'c'])


def select_data(l, dataset):
    df = None
    i = len(l)
    if i == 2:
        df = dataset[(dataset[l[0][0]] == l[0][1]) & (dataset[l[1][0]] == l[1][1])]

    if i == 3:
        df = dataset[(dataset[l[0][0]] == l[0][1]) & (dataset[l[1][0]] == l[1][1]) & (dataset[l[2][0]] == l[2][1])]

    return df


print(select_data(list_1, data))
print(select_data(list_2, data))     

必须有一种更简洁的方法来做到这一点。

Answer 1

使用 Dataframe.loc[<filtering_mask>, :] 进行简短过滤方法:

In [150]: data
Out[150]: 
   a  b  c
0  1  2  3
1  1  2  3
2  1  1  1
3  2  5  6

In [151]: def select_data(filter_lst, df):
     ...:     d = dict(filter_lst)
     ...:     res = df.loc[(df[list(d.keys())] == pd.Series(d)).all(axis=1)]
     ...:     return res
     ...: 
     ...: 

In [152]: select_data(list_1, data)
Out[152]: 
   a  b  c
3  2  5  6

In [153]: select_data(list_2, data)
Out[153]: 
   a  b  c
0  1  2  3
1  1  2  3