如何使用pandas的groupby功能计算含rho矩阵的嵌套双重求和？

你要计算的求和式本质是向量x关于矩阵rho的二次型x @ rho @ x，结合groupby和numpy的矩阵运算可以实现非常简洁的写法，完全不需要嵌套循环：

import pandas as pd
import numpy as np

# 示例数据
dftest=pd.DataFrame({
    "tau":[1,2,5,1,2,5],
    "x"  :[4,5,6,7,8,9],
    "cat":list("aaabbb")
})

rho = pd.DataFrame({
    1:[1, 0.9, 0.8],
    2:[0.9, 1, 0.7],
    5:[0.8, 0.7, 1]
}).T
rho = rho.rename(columns={0:1, 1:2, 2:5})

# 核心计算逻辑
def calc_quadratic(group):
    # 按tau排序保证和rho的行列顺序对齐
    sorted_group = group.sort_values("tau")
    x = sorted_group["x"].values
    # 提取当前分组tau对应的rho子矩阵，适配tau不全的场景
    tau_seq = sorted_group["tau"].tolist()
    rho_sub = rho.loc[tau_seq, tau_seq].values
    # 矩阵运算直接得到求和结果
    return x @ rho_sub @ x.T

# 分组计算
result = dftest.groupby("cat").apply(calc_quadratic).reset_index(name="rho_sum")

运行后得到的结果如下：

如果你的所有分组内的tau取值、顺序都和rho完全一致，可以省去提取子矩阵的步骤进一步提升性能：

def calc_quadratic_fast(group):
    x = group.sort_values("tau")["x"].values
    return x @ rho.values @ x.T

该方法的优势：

• 采用numpy向量化运算，性能远高于Python原生嵌套循环，数据量越大优势越明显
• 代码逻辑和数学表达式完全对应，可读性强
• 适配性好，就算分组内tau取值不全、顺序混乱也能正确计算

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