这是另一种方法。它更简洁、性能更高,并且具有columns 可以为空的优点(在这种情况下会返回整个数据帧)。
def filter(df, value, *columns):
return df.loc[df.loc[:, columns].eq(value).all(axis=1)]
说明
-
values = df.loc[:, columns] 只选择我们感兴趣的列。
-
masks = values.eq(value) 给出一个布尔数据框,表示与目标值相等。
-
mask = masks.all(axis=1) 跨列应用 AND(返回索引掩码)。请注意,您可以将 masks.any(axis=1) 用于 OR。
-
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)])
替代方案
对于少量列 (steven's answer 的解决方案比上述解决方案性能更高,但灵活性较差。照原样,它不适用于空的 columns 集,也无法在每列使用不同的值。
from operator import and_
def filter(df, value, *columns):
return df.loc[reduce(and_, (df[column] == value for column in columns))]
通过键检索 Series 对象 (df[column]) 比围绕列子集构造 DataFrame 对象 (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_, ...) 比内置的 Pandas all(axis=1) 慢得多。