基于多列计算分组滚动余额并实现值结转的技术求助

问题背景 #

原始DataFrame结构及数据：

import pandas as pd
import numpy as np

data = {'Date':['01/01/2022','01/02/2022','01/03/2022','01/04/2022','01/05/2022'],
        'ID': ['1', '1', '1', '1', '1'],
        'Sick': [0,0,8,0,4],
        'Grant': [80, np.nan, np.nan, np.nan, np.nan],
        'CarryForward': [80,80,80,80,72]}

df = pd.DataFrame(data)

输出结果：

Date ID  Sick  Grant  CarryForward
0  01/01/2022  1     0   80.0            80
1  01/02/2022  1     0    NaN            80
2  01/03/2022  1     8    NaN            80
3  01/04/2022  1     0    NaN            80
4  01/05/2022  1     4    NaN            72

需求说明 #

已按日期排序并按ID分组，需新增running_balance列，计算逻辑：

• 若Grant不为空，running_balance = Grant
• 若Grant为空且Sick为0/空，running_balance沿用前一日余额
• 若Grant为空且Sick不为0，running_balance = 前一日余额 - Sick

现有问题 #

原代码仅基于当日CarryForward计算，无法将扣减后的余额结转至下一行，导致第3行running_balance错误（应为72而非80）。

期望结果 #

Date ID  Sick  Grant  CarryForward  running_balance
0  01/01/2022  1     0   80.0            80               80
1  01/02/2022  1     0    NaN            80               80
2  01/03/2022  1     8    NaN            80               72
3  01/04/2022  1     0    NaN            80               72
4  01/05/2022  1     4    NaN            72               68

解决方案 #

方法1：高效累积计算（推荐） #

通过计算每日扣减额的累积值，结合初始余额生成滚动余额，避免逐行迭代，性能更优：

def compute_running_balance(group):
    # 获取初始余额：优先取第一个非空的Grant，若无则用首个CarryForward
    initial_balance = group['Grant'].dropna().iloc[0] if not group['Grant'].dropna().empty else group['CarryForward'].iloc[0]
    # 计算每日扣减额：Sick不为0时取对应值，否则为0
    daily_deduction = np.where(group['Sick'] != 0, group['Sick'], 0)
    # 生成累积扣减序列
    cumulative_deduction = daily_deduction.cumsum()
    # 计算滚动余额
    group['running_balance'] = initial_balance - cumulative_deduction
    return group

# 按ID分组后应用计算逻辑
df = df.groupby('ID', group_keys=False).apply(compute_running_balance)

方法2：逐行迭代计算 #

适合需要更复杂逻辑的场景，逐行处理并维护当前余额：

def calculate_running_balance(group):
    balance = []
    current_balance = None
    for _, row in group.iterrows():
        if pd.notnull(row['Grant']):
            current_balance = row['Grant']
        else:
            if row['Sick'] == 0 or pd.isnull(row['Sick']):
                # Sick为0/空时，沿用当前余额（首次则取CarryForward）
                if current_balance is None:
                    current_balance = row['CarryForward']
            else:
                # Sick不为0时，从当前余额扣减（首次则用CarryForward计算）
                if current_balance is None:
                    current_balance = row['CarryForward'] - row['Sick']
                else:
                    current_balance -= row['Sick']
        balance.append(current_balance)
    group['running_balance'] = balance
    return group

# 按ID分组应用逻辑
df = df.groupby('ID', group_keys=False).apply(calculate_running_balance)

两种方法均能得到符合期望的结果，且自动处理不同ID的分组隔离，确保余额计算互不干扰。

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