DataFrame按日期筛选col3前2高值并求和的实现方法

原始数据 #

先通过以下代码创建目标DataFrame：

import pandas as pd

d = {'Date': ['2020-1-1', '2020-1-2', '2020-1-3', '2020-1-1', '2020-1-2', 
              '2020-1-3','2020-1-1', '2020-1-2', '2020-1-3'], 
     'col2': ['A','A','A', 'B','B','B', 'C','C','C'],
     'col3':[0.01,0.02,0.03,0.02,0.03,0.04,0.05,0.1,0.01]}
df = pd.DataFrame(data=d)
df['Date'] = pd.to_datetime(df['Date'])

生成的DataFrame如下：

Date col2  col3
0 2020-01-01    A  0.01
1 2020-01-02    A  0.02
2 2020-01-03    A  0.03
3 2020-01-01    B  0.02
4 2020-01-02    B  0.03
5 2020-01-03    B  0.04
6 2020-01-01    C  0.05
7 2020-01-02    C  0.10
8 2020-01-03    C  0.01

需求说明 #

1. 按Date分组，每组取出col3值最高的2条记录（保留对应的col2信息）
2. 基于第一步结果，按日期对col3的值求和

实现步骤 #

1. 提取每组Top2记录 #

使用groupby结合nlargest方法，该方法针对分组场景做了优化，适合大数据量处理：

# 按Date分组，每组取col3最大的2条数据
top2_df = df.groupby('Date', group_keys=False).apply(lambda x: x.nlargest(2, 'col3')).reset_index(drop=True)

得到的中间结果（注：结果顺序与示例略有差异，但数据完全符合要求，若需指定顺序可额外添加排序逻辑）：

Date col2  col3
0 2020-01-01    C  0.05
1 2020-01-01    B  0.02
2 2020-01-02    C  0.10
3 2020-01-02    B  0.03
4 2020-01-03    B  0.04
5 2020-01-03    A  0.03

2. 按日期求和 #

基于上一步的结果，直接按Date分组对col3求和：

sum_df = top2_df.groupby('Date')['col3'].sum().reset_index(name='sum')

最终求和结果：

Date   sum
0 2020-01-01  0.07
1 2020-01-02  0.13
2 2020-01-03  0.07

方案优势 #

• nlargest比先排序再取前N的方式效率更高，尤其在数据量较大时表现明显
• 代码简洁易读，可直接适配更多col2分组和日期的场景

内容的提问来源于stack exchange，提问作者Beginner_01