编辑:我刚学会了一种更简洁的方法,使用 .transform
分组方法:
def get_max_rows(df):
B_maxes = df.groupby('A').B.transform(max)
return df[df.B == B_maxes]
B_maxes
是一个序列,与原始的
df
具有相同的索引,其中包含每个
A
组的
B
的最大值。您可以将许多函数传递给transform方法。我认为它们输出后要么是标量,要么是相同长度的向量。您甚至可以通过一些字符串作为常见函数名称(如
'median'
)来传递。这与Paul H的方法略有不同,因为'A'在结果中不会成为索引,但是您可以在之后轻松设置它。
import numpy as np
import pandas as pd
df_lots_groups = pd.DataFrame(np.random.rand(30000, 3), columns = list('BCD')
df_lots_groups['A'] = np.random.choice(range(10000), 30000)
%timeit get_max_rows(df_lots_groups)
100 loops, best of 3: 2.86 ms per loop
%timeit df_lots_groups.groupby('A').apply(lambda df: df[ df.B == df.B.max()])
1 loops, best of 3: 5.83 s per loop
编辑:
这里有一个抽象概念,允许您使用任何有效的比较运算符和任何有效的groupby方法从组中选择行:
def get_group_rows(df, group_col, condition_col, func=max, comparison='=='):
g = df.groupby(group_col)[condition_col]
condition_limit = g.transform(func)
df.query('condition_col {} @condition_limit'.format(comparison))
因此,例如,如果您想要在每个A组中选择所有高于中位数B值的行,则调用:
get_group_rows(df, 'A', 'B', 'median', '>')
一些例子:
%timeit get_group_rows(df_lots_small_groups, 'A', 'B', 'max', '==')
100 loops, best of 3: 2.84 ms per loop
%timeit get_group_rows(df_lots_small_groups, 'A', 'B', 'mean', '!=')
100 loops, best of 3: 2.97 ms per loop
df.query
和pd.eval
似乎非常适合这种用例。有关pd.eval()
函数系列的信息、特性和用例,请访问使用pd.eval()在pandas中进行动态表达式评估。 - cs95