基于日期的累积求和不符合预期：col3负值场景代码失效

需求说明 #

• 每日目标值为250，当col3（cum-daily_result）达到250后，对应col1分组的后续行col4需固定为首次达标时的col3值；若分组无达标行，col4直接等于col3。

原失效代码 #

原代码仅在col3首行为正值时输出正确，当col3首行为负值或分组全为负值时无法得到预期结果：

idx = df[df['col3'] >= 250].groupby('col1').head(1).index
df.loc[idx, 'col4'] = 1
df['col4'] = df.groupby('col1')['col4'].bfill() * df['col3']
df['col4'] = df.groupby('col1')['col4'].ffill().astype('int')

失效场景输入 #

import pandas as pd
data = {
    'col1': ['a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'c', 'c', 'c'],
    'col2': [-500, -250, -1000, -1392, 250, 250, 2238, 250, 2477, -3813, 250],
    'col3': [-500, -750, -1750, -1392, -1142, -892, 1346, 1596, 2477, -1336, -1086]
}
df = pd.DataFrame(data)

预期输出 #

col1  col2   col3  col4
0      a  -500   -500  -500
1      a  -250   -750  -750
2      a -1000  -1750 -1750
3      b -1392  -1392 -1392
4      b   250  -1142 -1142
5      b   250   -892  -892
6      b  2238   1346  1346
7      b   250   1596  1346
8      c  2477   2477  2477
9      c -3813  -1336  2477
10     c   250  -1086  2477

修正后的代码 #

import pandas as pd

def process_col4(df):
    # 按col1分组处理每个子数据集
    def handle_group(group):
        # 找到分组内第一个col3达标（>=250）的索引
        first_qualified_idx = group[group['col3'] >= 250].index.min()
        
        if pd.notna(first_qualified_idx):
            # 达标前的行col4等于对应col3值
            group.loc[:first_qualified_idx, 'col4'] = group.loc[:first_qualified_idx, 'col3']
            # 达标后的行固定为首次达标时的col3值
            qualified_value = group.loc[first_qualified_idx, 'col3']
            group.loc[first_qualified_idx+1:, 'col4'] = qualified_value
        else:
            # 无达标行时，col4直接等于col3
            group['col4'] = group['col3']
        
        return group
    
    # 应用分组逻辑并转换为整数类型
    df = df.groupby('col1', group_keys=False).apply(handle_group)
    df['col4'] = df['col4'].astype(int)
    return df

# 测试执行
data = {
    'col1': ['a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'c', 'c', 'c'],
    'col2': [-500, -250, -1000, -1392, 250, 250, 2238, 250, 2477, -3813, 250],
    'col3': [-500, -750, -1750, -1392, -1142, -892, 1346, 1596, 2477, -1336, -1086]
}
df = pd.DataFrame(data)
df = process_col4(df)
print(df)

代码说明 #

• 针对每个col1分组单独处理，避免跨分组干扰
• 精准定位首次达标位置，区分达标前后的赋值逻辑
• 兼容无达标行的分组场景，直接映射col3值到col4
• 统一转换为整数类型，匹配预期输出格式

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