如何用Pandas GroupBy实现DataFrame的复杂双重求和计算

问题背景 #

给定如下结构的Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    'Race_ID': [1,1,1,1,2,2,2],
    'Date': ['1/1/2023','1/1/2023','1/1/2023','1/1/2023','11/9/2022','11/9/2022','11/9/2022'],
    'Student_ID': [1,2,3,4,1,2,3],
    'a': [3,2,1,4,2,3,1],
    'b': [1,2,4,1,6,2,6]
})

需要按Race_ID分组，为每行新增sum列，其值由以下双重求和公式计算：
$$\sum_{1\leq k \leq n_{RaceID} \atop{k\not=i}}\sum_{1\leq j \leq n_{RaceID} \atop{j \notin {i, k}}}a_kb_j$$
其中：

• $i$为当前行在分组内的索引
• $n_{RaceID}$为对应分组的行数
• 实际场景中求和项并非可分离的a_kb_j形式，因此无法通过数学拆分简化计算

解决方案 #

方法1：分组遍历实现（直观通用） #

通过groupby.apply对每个分组单独处理，遍历每行$i$，再遍历符合条件的$k$和$j$计算求和：

def calculate_group_sum(group):
    n = len(group)
    a_vals = group['a'].values
    b_vals = group['b'].values
    sum_results = []
    
    for i in range(n):
        total = 0
        # 遍历所有k≠i
        for k in range(n):
            if k == i:
                continue
            # 遍历所有j∉{i,k}
            for j in range(n):
                if j == i or j == k:
                    continue
                # 替换此处为实际的自定义函数，示例为a_k*b_j
                total += a_vals[k] * b_vals[j]
        sum_results.append(total)
    
    group['sum'] = sum_results
    return group

# 应用到分组数据
result_df = df.groupby('Race_ID', group_keys=False).apply(calculate_group_sum)

方法2：向量化优化（适合大数据量） #

利用Numpy生成掩码矩阵，减少循环次数，提升效率：

import numpy as np

def calculate_group_sum_vectorized(group):
    n = len(group)
    a = group['a'].values.reshape(-1, 1)
    b = group['b'].values.reshape(1, -1)
    
    # 生成求和项矩阵，替换此处为实际自定义函数的矩阵计算
    # 示例为a_k*b_j的矩阵
    term_matrix = a @ b
    
    sum_results = []
    for i in range(n):
        # 生成掩码：排除j=i、k=i、j=k的情况
        mask = np.ones((n, n), dtype=bool)
        mask[i, :] = False  # 排除j=i
        mask[:, i] = False  # 排除k=i
        np.fill_diagonal(mask, False)  # 排除j=k
        
        # 对符合条件的项求和
        total = term_matrix[mask].sum()
        sum_results.append(total)
    
    group['sum'] = sum_results
    return group

# 应用到分组数据
result_df = df.groupby('Race_ID', group_keys=False).apply(calculate_group_sum_vectorized)

验证结果 #

运行上述代码后，result_df将生成符合预期的输出：

Race_ID        Date  Student_ID  a  b  sum
0        1    1/1/2023           1  3  1   37
1        1    1/1/2023           2  2  2   37
2        1    1/1/2023           3  1  4   25
3        1    1/1/2023           4  4  1   31
4        2  11/9/2022           1  2  6   20
5        2  11/9/2022           2  3  2   18
6        2  11/9/2022           3  1  6   22

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