是否可以使用 pandas 和 groupby 计算嵌套总和？

Question

假设我有一个包含两个数字列和一个分类列的数据框：

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

输出：

我可以轻松计算每个 cat 的 x 的平方和：

dftest.groupby(["cat"]).apply(
    lambda s: pd.Series({
        "sum_x^2":(s["x"]**2).sum()
    })
)

但现在我的问题是：对于一些具有正确维度和 row/column 索引的方阵 rho，我如何计算 Sum_{i,j} rho(tau_i, tau_j) * x_i * x_j 而不是 Sum_i (x_i^2)？作为一个具体的例子，假设我有以下 rho:

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})

对于类别a，我想计算：

rho[1,1] * x[1] * x[1] + rho[1,2] * x[1] * x[2] + rho[1,5] * x[1] * x[5] + ...
... rho[2,1] * x[2] * x[1] + ...

或按照给出的示例：

1*4*4 + 0.9*4*5 + 0.8*4*6 + ...
0.9*5*4 + ...

目前我正在通过循环遍历 groupby 对象并循环遍历 rho 中的索引值来执行此操作，但我想知道我是否缺少更优雅的方法。

Answer 1

首先按组 dftest 对列进行透视 1,2,5:

df = dftest.pivot('cat','tau','x')
print (df)
tau  1  2  5
cat         
a    4  5  6
b    7  8  9

然后由 DataFrame.stack 重塑 rho 为 MultiIndex Series:

s = rho.stack()
print (s)
1  1    1.0
   2    0.9
   5    0.8
2  1    0.9
   2    1.0
   5    0.7
5  1    0.8
   2    0.7
   5    1.0
dtype: float64

DataFrame.reindex 在不同级别的重复值：

df1 = df.reindex(s.index, level=0, axis=1)
print (df1)
     1        2        5      
     1  2  5  1  2  5  1  2  5
cat                           
a    4  4  4  5  5  5  6  6  6
b    7  7  7  8  8  8  9  9  9

df2 = df.reindex(s.index, level=1, axis=1)
print (df2)
     1        2        5      
     1  2  5  1  2  5  1  2  5
cat                           
a    4  5  6  4  5  6  4  5  6
b    7  8  9  7  8  9  7  8  9

最后一个是多个sum:

out = df1.mul(df2).mul(s).sum(axis=1)
print (out)
cat
a    193.4
b    496.4
dtype: float64

是否可以使用 pandas 和 groupby 计算嵌套总和？

Is it possible to compute nested sums using pandas and groupby?

pandas

pandas-groupby