如何高效识别DataFrame中事务的冲销记录？

问题背景 #

我有一组包含MatterNum（事务编号）、WorkDate（工作日期）、Amount（金额）的事务数据。录入人员常因失误需通过录入负数金额冲销错误成本，需要按事务编号和工作日期分组，对比金额识别冲销记录及被冲销记录。

数据示例 #

期望输出 #

新增Reversal?列标记冲销记录：

原实现的性能瓶颈 #

原代码用apply逐行处理，每次调用函数都要重复筛选分组数据，百万级数据下速度极慢：

import pandas as pd 

data = [
        [1,'1/2/2022',10],
        [1,'1/2/2022',15],
        [1,'1/2/2022',-10],
        [2,'1/4/2022',12],
        [2,'1/5/2022',-5],
        [2,'1/5/2022',5]
    ]
df = pd.DataFrame(data, columns=['MatterNum','WorkDate','Amount'])

def rev_check(MatterNum, workDate, WorkAmt, df):
    funcDF = df.loc[(df['MatterNum'] == MatterNum) & (df['WorkDate'] == workDate)] 

    listCheck = funcDF['Amount'].tolist()
    if WorkAmt*-1 in listCheck:
        return 'yes'

df['reversal?'] = df.apply(lambda row: rev_check(row.MatterNum, row.WorkDate, row.Amount, df), axis=1)

优化方案 #

核心思路是减少重复计算，用分组统计+向量化操作替代逐行循环，把时间复杂度从O(n²)降到O(n)：

方案一：分组生成金额集合后匹配 #

import pandas as pd 

data = [
        [1,'1/2/2022',10],
        [1,'1/2/2022',15],
        [1,'1/2/2022',-10],
        [2,'1/4/2022',12],
        [2,'1/5/2022',-5],
        [2,'1/5/2022',5]
    ]
df = pd.DataFrame(data, columns=['MatterNum','WorkDate','Amount'])

# 按分组生成该组的金额集合，只执行一次分组统计
grouped_amounts = df.groupby(['MatterNum', 'WorkDate'])['Amount'].agg(set).reset_index(name='AmountSet')

# 把金额集合合并回原表
df = df.merge(grouped_amounts, on=['MatterNum', 'WorkDate'], how='left')

# 批量判断每个金额的相反数是否在同组集合中
df['Reversal?'] = df.apply(lambda x: 'yes' if (-x['Amount'] in x['AmountSet']) else 'no', axis=1)

# 清理临时列
df.drop('AmountSet', axis=1, inplace=True)

print(df)

方案二：用transform完全避免额外合并 #

进一步优化，直接通过groupby.transform在分组内完成判断，代码更简洁：

import pandas as pd 

data = [
        [1,'1/2/2022',10],
        [1,'1/2/2022',15],
        [1,'1/2/2022',-10],
        [2,'1/4/2022',12],
        [2,'1/5/2022',-5],
        [2,'1/5/2022',5]
    ]
df = pd.DataFrame(data, columns=['MatterNum','WorkDate','Amount'])

# 定义分组内的判断逻辑：先取该组金额集合，再逐个判断相反数是否存在
def check_reversal(amount_series):
    amount_set = set(amount_series)
    return amount_series.apply(lambda x: 'yes' if (-x in amount_set) else 'no')

# 分组后直接生成Reversal?列
df['Reversal?'] = df.groupby(['MatterNum', 'WorkDate'])['Amount'].transform(check_reversal)

print(df)

性能提升说明 #

• 原代码：百万行数据需数分钟，因为每次循环都要重新筛选分组数据，重复计算量大；
• 优化后代码：百万行数据仅需数秒，分组统计只执行一次，后续都是批量的集合查找（O(1)时间复杂度）。

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