如何按id筛选满足day条件且value总和≤20的DataFrame行？

步骤说明 #

针对每个id，我们分两部分筛选行：

1. 直接保留day > day_before的行；
2. 对day ≤ day_before的行，按day降序排序后从大到小累加value，仅保留累计和≤20的行。

完整代码 #

import pandas as pd

# 初始化原始数据
df = pd.DataFrame({
    'id': [1,1,1,1,1,2,2,2,2],
    'day': [1,2,3,4,5,1,2,3,4],
    'value': [5,6,7,8,9,11,12,13,14],
    'day_before': [4,4,4,4,4,3,3,3,3]
})

def filter_group(group):
    # 保留day大于day_before的行
    keep_part1 = group[group['day'] > group['day_before']]
    # 提取day小于等于day_before的子组
    part2 = group[group['day'] <= group['day_before']]
    
    if part2.empty:
        return keep_part1
    
    # 按day降序排序，从最大的day开始累加
    part2_sorted = part2.sort_values('day', ascending=False)
    # 计算value的累计和
    part2_sorted['cum_sum'] = part2_sorted['value'].cumsum()
    # 筛选累计和≤20的行，移除辅助列
    keep_part2 = part2_sorted[part2_sorted['cum_sum'] <= 20].drop('cum_sum', axis=1)
    
    # 合并两部分结果，按day排序并重置索引
    return pd.concat([keep_part1, keep_part2]).sort_values('day').reset_index(drop=True)

# 按id分组应用筛选逻辑
final_df = df.groupby('id', group_keys=False).apply(filter_group)

# 输出结果
print(final_df)

输出结果 #

id  day  value  day_before
0   1    3      7           4
1   1    4      8           4
2   1    5      9           4
3   2    3     13           3
4   2    4     14           3

逻辑解释 #

• 分组处理：按id拆分数据，确保每个id的筛选独立进行；
• 第一部分筛选：day > day_before的行无累加限制，直接保留；
• 第二部分筛选：将day ≤ day_before的行按day从大到小排序，累加value后仅保留累计和不超过20的行——这样能保证从最大的符合条件的day开始，尽可能多地保留行且不触发总和限制；
• 合并结果：将两部分有效行合并，按day排序后得到最终数据。

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