pandas:是否可以使用任意长的布尔标准过滤数据帧？

Question

如果您确切地知道如何过滤数据帧,那么解决方案很简单:

df[(df.A == 1) & (df.B == 1)]

但是,如果您接受用户输入并且事先不知道用户想要使用多少条件,该怎么办？例如,用户想要一个过滤的数据框,其中列[A,B,C] == 1.是否可以执行以下操作:

def filterIt(*args, value):
    return df[(df.*args == value)]

所以如果用户调用filterIt(A, B, C, value=1),则返回:

df[(df.A == 1) & (df.B == 1) & (df.C == 1)]

Answer 1

我认为最优雅的方法是使用df.query(),在这里您可以构建一个包含所有条件的字符串,例如:

import pandas as pd
import numpy as np

cols = {}
for col in ('A', 'B', 'C', 'D', 'E'):
    cols[col] = np.random.randint(1, 5, 20)
df = pd.DataFrame(cols)

def filter_df(df, filter_cols, value):
    conditions = []
    for col in filter_cols:
        conditions.append('{c} == {v}'.format(c=col, v=value))
    query_expr = ' and '.join(conditions)
    print('querying with: {q}'.format(q=query_expr))
    return df.query(query_expr)

示例输出(由于随机生成的数据,您的结果可能会有所不同):

filter_df(df, ['A', 'B'], 1)
querying with: A == 1 and B == 1
    A  B  C  D  E
6   1  1  1  2  1
11  1  1  2  3  4

Answer 2

这是另一种方法.它更干净,性能更高,并且具有columns可以为空的优点(在这种情况下返回整个数据框).

def filter(df, value, *columns):
    return df.loc[df.loc[:, columns].eq(value).all(axis=1)]

说明

1. values = df.loc[:, columns] 只选择我们感兴趣的列.
2. masks = values.eq(value) 给出一个布尔数据框,表示与目标值的相等性.
3. mask = masks.all(axis=1)在列之间应用AND(返回索引掩码).请注意,您可以使用masks.any(axis=1)OR.
4. return df.loc[mask] 将索引掩码应用于数据框.

演示

import numpy as np
import pandas as pd

df = pd.DataFrame(np.random.randint(0, 2, (100, 3)), columns=list('ABC'))

# both columns
assert np.all(filter(df, 1, 'A', 'B') == df[(df.A == 1) & (df.B == 1)])

# no columns
assert np.all(filter(df, 1) == df)

# different values per column
assert np.all(filter(df, [1, 0], 'A', 'B') == df[(df.A == 1) & (df.B == 0)])

替代

对于少数列(<5),基于史蒂文的答案,以下解决方案比上述解决方案更具性能,但灵活性较差.原样,它不适用于空columns集,并且不会使用每列不同的值.

from operator import and_

def filter(df, value, *columns):
    return df.loc[reduce(and_, (df[column] == value for column in columns))]

Series通过key(df[column])检索对象比DataFrame围绕column(df.loc[:, columns])子集构造对象要快得多.

In [4]: %timeit df['A'] == 1
100 loops, best of 3: 17.3 ms per loop

In [5]: %timeit df.loc[:, ['A']] == 1
10 loops, best of 3: 48.6 ms per loop

然而,当处理更多列时,这种加速变得可以忽略不计.瓶颈变成了将面具加在一起,这reduce(and_, ...)比熊猫内置的慢得多all(axis=1).