如何优化Pandas DataFrame中找和为0多行组合的算法性能？

问题背景 #

我有一个包含10000行的Pandas DataFrame，需要找出特定列（Amount）数值之和为0的多行组合（包括3行、4行、5行组合）。

示例DataFrame #

ID_Key          Amount
10               12.4
12               -26.6
13                14.2        
14                15
17                 4.5
18                -9
19                94
20                -6

期望结果 #

Combinations        Sum
(10,12,13)            0
(14,18,20)            0

当前使用的3行组合求和代码 #

from itertools import combinations
lst = [] 
t_counter=0
#all combinations ID_key consisting of length 3
for tuple_nums in set(combinations(df['ID_Key'], 3)):

    if df.shape[0]>2:
        t_counter=t_counter+1
        if df.loc[df['ID_Key'].isin(tuple_nums)].empty==False:
            if df.loc[df['ID_Key'].isin(tuple_nums), 'Amount'].sum()==0:
                lst.append([tuple_nums,df.loc[df['ID_Key'].isin(tuple_nums), 'Amount'].sum()])
                df=df.loc[~df['ID_Key'].isin(tuple_nums)]


    else:
        break



df_final=pd.DataFrame(lst, columns=['Combinations', 'Sum'])

该算法在数据量超过30行时就会变得很慢，即使每次找到符合条件的组合后缩小DataFrame规模，仍需遍历所有可能的组合。请问如何降低算法的时间复杂度，以处理10000行的数据？

优化方案 #

直接遍历所有组合的时间复杂度为O(C(n,k))，当n=10000时完全不可行，必须改用哈希表存储中间结果的思路，减少重复计算。

1. 3行组合的核心优化 #

3行和为0等价于 a + b + c = 0 → a + b = -c。先预存所有两数之和及其对应的ID对，再遍历每个元素查找补数：

import pandas as pd
from collections import defaultdict

def find_3sum_groups(df):
    # 构建ID与Amount的映射字典
    id_amount_map = df.set_index('ID_Key')['Amount'].to_dict()
    id_list = list(id_amount_map.keys())
    
    # 存储两数之和对应的ID对（避免重复存储同一组合）
    sum_pair_map = defaultdict(list)
    for i in range(len(id_list)):
        for j in range(i+1, len(id_list)):
            id_a, id_b = id_list[i], id_list[j]
            sum_ab = id_amount_map[id_a] + id_amount_map[id_b]
            sum_pair_map[sum_ab].append((id_a, id_b))
    
    result = []
    used_ids = set()  # 标记已使用的ID，避免重复组合
    
    for id_c in id_list:
        if id_c in used_ids:
            continue
        # 计算需要匹配的补数
        target = -id_amount_map[id_c]
        if target in sum_pair_map:
            for (id_a, id_b) in sum_pair_map[target]:
                # 确保三个ID不重复且未被使用
                if id_c not in (id_a, id_b) and id_a not in used_ids and id_b not in used_ids:
                    # 对ID排序，避免重复组合（如(10,12,13)和(12,10,13)视为同一组合）
                    combo = tuple(sorted((id_a, id_b, id_c)))
                    result.append({'Combinations': combo, 'Sum': 0.0})
                    used_ids.update(combo)
    
    return pd.DataFrame(result)

2. 扩展到4行、5行组合 #

• 4行组合：可拆分为 a+b = -(c+d)，先预存所有两数之和的组合，再查找和为0的两对；或拆分为 a+b+c = -d，复用3行组合的结果再匹配单个元素。
• 5行组合：可拆分为 a+b+c+d = -e，先预存所有4数之和的组合（建议先过滤绝对值过大的元素减少计算量）；或拆分为 (a+b)+(c+d) = -e，利用两数之和的中间结果。

3. 通用优化细节 #

• 浮点数精度处理：不要直接判断sum == 0，改用abs(sum) < 1e-9的阈值判断，避免精度误差导致漏判。
• 去重预处理：将相同Amount的ID分组，比如多个ID的Amount都是10，可合并处理，减少计算量。
• 剪枝操作：提前过滤绝对值过大的元素——如果某个元素的绝对值大于其他所有元素的和，它不可能出现在任何和为0的组合中。
• 分批处理：数据量过大时，可分成多个小批次，先处理批次内的组合，再跨批次匹配（需注意避免重复组合）。

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