寻求Pandas条件求和高效实现方案的技术建议

我是Python和Pandas新手，现有一段计算特定条件下vesting_value_CAD总和的代码，性能较差，尝试过cumsum但结果不符合预期。

需求说明 #

为每个trancheID计算vesting_value_CAD的总和，求和范围需满足：

• 与当前行拥有相同employeeID、groupName、vesting_year
• agreementDate ≤ 当前行agreementDate
• 排除当前行

原始代码 #

import pandas as pd
from datetime import datetime

# 构建数据
data = {
    'employeeID': [2, 2, 2, 2, 2, 2, 2],
    'groupName': ['A', 'A', 'A', 'A', 'A', 'B', 'A'],
    'agreementID': [7, 7, 1, 1, 8, 9, 6],
    'agreementDate': ['3/1/2025', '3/1/2025', '4/1/2025', '3/1/2025', '2/1/2025', '3/1/2025', '3/1/2025'],
    'trancheID': [28, 29, 26, 27, 30, 31, 32],
    'vesting_year': [2025, 2026, 2026, 2027, 2026, 2026, 2026],
    'vesting_value_CAD': [200, 300, 400, 500, 50, 30, 40]
}

df = pd.DataFrame(data)

# 转换日期格式
df['agreementDate'] = pd.to_datetime(df['agreementDate'], format='%m/%d/%Y')

# 逐行计算符合条件的总和
def calculate_total_vesting_value(row):
    filtered_df = df[(df['employeeID'] == row['employeeID']) & 
                     (df['groupName'] == row['groupName']) & 
                     (df['vesting_year'] == row['vesting_year']) & 
                     (df['agreementDate'] <= row['agreementDate']) & 
                     (df['trancheID'] != row['trancheID'])]
    return filtered_df['vesting_value_CAD'].sum()

df['total_vesting_value_CAD'] = df.apply(calculate_total_vesting_value, axis=1)

print(df)

优化方案 #

原始代码用df.apply逐行循环，时间复杂度为O(n²)，数据量大时性能极差。以下是基于Pandas向量化操作的优化方案，时间复杂度为O(n log n)（主要来自排序）：

import pandas as pd
from datetime import datetime

# 构建数据
data = {
    'employeeID': [2, 2, 2, 2, 2, 2, 2],
    'groupName': ['A', 'A', 'A', 'A', 'A', 'B', 'A'],
    'agreementID': [7, 7, 1, 1, 8, 9, 6],
    'agreementDate': ['3/1/2025', '3/1/2025', '4/1/2025', '3/1/2025', '2/1/2025', '3/1/2025', '3/1/2025'],
    'trancheID': [28, 29, 26, 27, 30, 31, 32],
    'vesting_year': [2025, 2026, 2026, 2027, 2026, 2026, 2026],
    'vesting_value_CAD': [200, 300, 400, 500, 50, 30, 40]
}

df = pd.DataFrame(data)
df['agreementDate'] = pd.to_datetime(df['agreementDate'], format='%m/%d/%Y')

# 分组键定义
group_keys = ['employeeID', 'groupName', 'vesting_year']

# 步骤1：计算每个(分组+日期)组合的总价值
df['date_group_sum'] = df.groupby(group_keys + ['agreementDate'])['vesting_value_CAD'].transform('sum')

# 步骤2：按分组键和日期排序，保证累计和顺序正确
df_sorted = df.sort_values(group_keys + ['agreementDate'])

# 步骤3：计算分组内到当前日期的累计总价值
df_sorted['total_to_date'] = df_sorted.groupby(group_keys)['date_group_sum'].cumsum()

# 步骤4：减去当前行价值，得到排除自身后的总和
df_sorted['total_vesting_value_CAD'] = df_sorted['total_to_date'] - df_sorted['vesting_value_CAD']

# 恢复原始数据顺序（可选）
df = df_sorted.sort_index()

print(df)

优化逻辑说明 #

1. 日期分组求和：先计算同一分组内、同一日期的所有vesting_value_CAD总和，确保同日期的行能互相计入对方的结果。
2. 排序：按分组键和日期排序，保证累计和按时间顺序计算。
3. 累计求和：对每个分组计算到当前日期的累计总价值，包含所有≤当前日期的行的价值。
4. 排除当前行：用累计总价值减去当前行的vesting_value_CAD，得到符合要求的结果。

该方案完全使用Pandas向量化操作，避免逐行循环，数据量越大性能提升越明显。

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