在调用df.groupby(...).apply(foo) 时,foo 返回的对象类型会影响结果融合在一起的方式。
如果返回一个 Series,则 Series 的索引成为最终结果的列,而 groupby 键成为索引(有点绕口令)。
如果改为返回一个DataFrame,最终结果使用DataFrame的索引作为索引值,DataFrame的列作为列(非常明智)。
因此,您可以通过将 Series 转换为 DataFrame 来安排所需的输出类型。
在 Pandas 0.13 中,您可以使用 to_frame().T 方法:
def maxrow(x, col):
return x.loc[x[col].argmax()].to_frame().T
result = df.groupby('c1').apply(maxrow, 'c3')
result = result.reset_index(level=0, drop=True)
print(result)
产量
c1 c2 c3
1 a c 3
4 b c 12
在 Pandas 0.12 或更早版本中,相当于:
def maxrow(x, col):
ser = x.loc[x[col].idxmax()]
df = pd.DataFrame({ser.name: ser}).T
return df
顺便说一下,behzad.nouri's clever and elegant solution 对于小型 DataFrame 来说比我的要快。
sort 将时间复杂度从 O(n) 提升到 O(n log n),因此当应用于更大的 DataFrame 时,它比上面显示的 to_frame 解决方案要慢。
这是我对其进行基准测试的方法:
import pandas as pd
import numpy as np
import timeit
def reset_df_first(df):
df2 = df.reset_index()
result = df2.groupby('c1').apply(lambda x: x.loc[x['c3'].idxmax()])
result.set_index(['index'], inplace=True)
return result
def maxrow(x, col):
result = x.loc[x[col].argmax()].to_frame().T
return result
def using_to_frame(df):
result = df.groupby('c1').apply(maxrow, 'c3')
result.reset_index(level=0, drop=True, inplace=True)
return result
def using_sort(df):
return df.sort('c3').groupby('c1', as_index=False).tail(1)
for N in (100, 1000, 2000):
df = pd.DataFrame({'c1': {0: 'a', 1: 'a', 2: 'a', 3: 'b', 4: 'b', 5: 'b'},
'c2': {0: 'a', 1: 'c', 2: 'b', 3: 'b', 4: 'c', 5: 'a'},
'c3': {0: 1, 1: 3, 2: 2, 3: 10, 4: 12, 5: 7}})
df = pd.concat([df]*N)
df.reset_index(inplace=True, drop=True)
timing = dict()
for func in (reset_df_first, using_to_frame, using_sort):
timing[func] = timeit.timeit('m.{}(m.df)'.format(func.__name__),
'import __main__ as m ',
number=10)
print('For N = {}'.format(N))
for func in sorted(timing, key=timing.get):
print('{:<20}: {:<0.3g}'.format(func.__name__, timing[func]))
print
产量
For N = 100
using_sort : 0.018
using_to_frame : 0.0265
reset_df_first : 0.0303
For N = 1000
using_to_frame : 0.0358 \
using_sort : 0.036 / this is roughly where the two methods cross over in terms of performance
reset_df_first : 0.0432
For N = 2000
using_to_frame : 0.0457
reset_df_first : 0.0523
using_sort : 0.0569
(reset_df_first 是我尝试过的另一种可能性。)