如何高效实现Pandas分组列的非自乘两两乘积求和

原始数据集 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    's1': [1,1,2,2,2,2,3,3,3,3],
    's2': [2,3,1,3,3,3,4,4,5,5],
    'l': [1,1,1,1,2,3,1,2,1,2],
    'X_ikl': [0,0,1,8,2,4,1,1,2,10]
})

需求说明 #

按(s1,s2)列分组，对每组内的X_ikl进行不包含自身的两两相乘，将所有乘积结果求和；若组内仅1条数据，结果为0。要求高效实现（避免for循环）。

示例预期结果 #

(1,2): 0
(1,3): 0
(2,1): 0
(2,3): 8*2 + 8*4 + 2*4 = 56
(3,4): 1*1 = 1
(3,5): 2*10 = 20

用户尝试的错误代码 #

df.groupby(["s1","s2"]).apply(lambda x : np.sum(x["X_ikl"]*x["X_ikl"]))

该代码计算的是组内元素自乘后的平方和，与需求不符。

高效解法 #

思路：用数学公式简化计算 #

对于一组数[a1,a2,...,an]，所有两两不同元素的乘积之和可以通过以下公式推导得出：
$$\frac{(\sum_{i=1}^n a_i)^2 - \sum_{i=1}^n a_i^2}{2}$$
推导逻辑：
$$(\sum a_i)^2 = \sum a_i^2 + 2\sum_{i<j}a_i a_j$$
移项后直接得到两两乘积之和的表达式，且当组内只有1个元素时，分子为0，结果自动为0，完全符合需求。

代码实现 #

# 分组计算总和、平方和
grouped = df.groupby(["s1", "s2"])["X_ikl"].agg(
    sum_total="sum",
    sum_squares=lambda x: (x**2).sum()
)

# 计算两两乘积之和
grouped["pairwise_product_sum"] = (grouped["sum_total"] ** 2 - grouped["sum_squares"]) // 2

# 输出结果
print(grouped["pairwise_product_sum"])

最终输出 #

s1  s2
1   2     0
    3     0
2   1     0
    3    56
3   4     1
    5    20
Name: pairwise_product_sum, dtype: int64

该方案全程使用Pandas内置分组聚合操作，无循环，处理大数据集效率极高，结果完全匹配预期。

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