如何按trade_date分组筛选Pandas DataFrame中return列前20%的行？

问题描述 #

给定如下Pandas DataFrame：

import pandas as pd

data = {'trade_date': {1350: 20151201,
  6175: 20151201,
  3100: 20151201,
  5650: 20151201,
  3575: 20151201,
     1: 20170301,
     2: 20170301},
 'comId': {1350: '257762',
  6175: '1038328',
  3100: '315476',
  5650: '658776',
  3575: '329376',
     1: '123456',
     2: '987654'},
 'return': {1350: -0.0018,
  6175: 0.0023,
  3100: -0.0413,
  5650: 0.1266,
  3575: 0.0221,
  1: '0.9',
  2: '0.01'}}

df = pd.DataFrame(data)

需要完成的筛选逻辑：
针对每个trade_date值，保留return列值排名前20%的条目，条目数量按四舍五入取整（例如某日期关联9家公司，20%×9=1.8，则保留前2条）。

期望输出结果：

{'trade_date': {5650: 20151201, 1: 20170301},
 'comId': {5650: '658776', 1: '123456'},
 'return': {5650: 0.1266, 1: 0.9}}

实现方案 #

步骤说明 #

1. 统一数据类型：return列同时存在数值和字符串类型，需先转为数值类型才能正确排序。
2. 分组筛选：按trade_date分组后，计算每组需保留的条目数（总条数×20%后四舍五入），再提取每组中return值最高的前N条。

完整代码 #

import pandas as pd

# 原始数据初始化
data = {'trade_date': {1350: 20151201,
  6175: 20151201,
  3100: 20151201,
  5650: 20151201,
  3575: 20151201,
     1: 20170301,
     2: 20170301},
 'comId': {1350: '257762',
  6175: '1038328',
  3100: '315476',
  5650: '658776',
  3575: '329376',
     1: '123456',
     2: '987654'},
 'return': {1350: -0.0018,
  6175: 0.0023,
  3100: -0.0413,
  5650: 0.1266,
  3575: 0.0221,
  1: '0.9',
  2: '0.01'}}

df = pd.DataFrame(data)

# 1. 将return列统一转换为数值类型
df['return'] = pd.to_numeric(df['return'])

# 2. 定义分组筛选函数
def filter_top_percent(group):
    # 计算需保留的条目数，四舍五入处理
    keep_count = round(len(group) * 0.2)
    # 若计算结果为0，强制保留1条（匹配示例中2条数据取1条的逻辑）
    keep_count = max(keep_count, 1)
    # 提取return值最大的前N条
    return group.nlargest(keep_count, 'return')

# 分组应用筛选逻辑
result_df = df.groupby('trade_date', group_keys=False).apply(filter_top_percent)

# 输出结果（转为字典格式匹配期望输出）
print(result_df.to_dict())

代码解释 #

• pd.to_numeric()：解决return列混合类型问题，确保排序逻辑正确。
• groupby().apply()：按日期分组后，对每个分组执行自定义筛选逻辑。
• nlargest()：高效提取每组中指定列的前N个最大值条目，比先排序再切片更高效。
• max(keep_count, 1)：处理四舍五入后为0的场景（如2条数据×20%=0.4，四舍五入为0，此时强制保留1条，符合示例输出）。

内容的提问来源于stack exchange，提问作者A.N.