如何在Python中用groupby+lambda正确实现组间/组内平方和矩阵计算

本人熟悉R语言，刚接触Python。以经典iris数据集为例，先在R中完成数据预处理与组统计计算，结果与手动计算一致。现尝试将R逻辑迁移至Python，希望通过groupby结合lambda函数的写法实现，避免拆分数据再整合的繁琐方式。已编写Python代码，前序步骤均正常，但最后计算组间平方和矩阵H与组内平方和矩阵E的代码输出错误，寻求解决方法。

R语言实现代码 #

str(iris)
## 'data.frame':   150 obs. of  5 variables:
##  $ Sepal.Length: num  5.1 4.9 4.7 4.6 5 5.4 4.6 5 4.4 4.9 ...
##  $ Sepal.Width : num  3.5 3 3.2 3.1 3.6 3.9 3.4 3.4 2.9 3.1 ...
##  $ Petal.Length: num  1.4 1.4 1.3 1.5 1.4 1.7 1.4 1.5 1.4 1.5 ...
##  $ Petal.Width : num  0.2 0.2 0.2 0.2 0.2 0.4 0.3 0.2 0.2 0.1 ...
##  $ Species     : Factor w/ 3 levels "setosa","versicolor",..: 1 1 1 1 1 1 1 1 1 1 ...

dat = iris  #iris dataset in base R
names(dat) = c(paste0('x',1:4),'g') #rename columns
dat$g = as.numeric(dat$g) #re-coding 3-level factor 
dat[,1:4] = dat[,1:4]*10 #enlarge x1-x4 values; just a test!
dim(dat)
## [1] 150   5
head(dat) #g=group, 3-level: 1,2,3
##   x1 x2 x3 x4 g
## 1 51 35 14  2 1
## 2 49 30 14  2 1
## 3 47 32 13  2 1
## 4 46 31 15  2 1
## 5 50 36 14  2 1
## 6 54 39 17  4 1

by(dat[,-5],dat[,5],nrow) #row size for each group; n_i
## dat[, 5]: 1
## [1] 50
## ------------------------------------------------------------ 
## dat[, 5]: 2
## [1] 50
## ------------------------------------------------------------ 
## dat[, 5]: 3
## [1] 50

by(dat[,-5],dat[,5],colMeans) #mean vector for data in each group; \bar{x}_i
## dat[, 5]: 1
##    x1    x2    x3    x4 
## 50.06 34.28 14.62  2.46 
## ------------------------------------------------------------ 
## dat[, 5]: 2
##    x1    x2    x3    x4 
## 59.36 27.70 42.60 13.26 
## ------------------------------------------------------------ 
## dat[, 5]: 3
##    x1    x2    x3    x4 
## 65.88 29.74 55.52 20.26 

m = dat[,-5] |> colMeans(); m  #mean vector for whole data; \bar{x}
##      x1       x2       x3       x4 
## 58.43333 30.57333 37.58000 11.99333 

# by(dat[,-5],dat[,5],cov) #covariance matrix for data in each group

H = by(dat[,-5],dat[,5],function(x) nrow(x)*tcrossprod(colMeans(x)-m) ) |> Reduce(f="+") ; H
##           [,1]      [,2]     [,3]      [,4]
## [1,]  6321.213 -1995.267 16524.84  7127.933
## [2,] -1995.267  1134.493 -5723.96 -2293.267
## [3,] 16524.840 -5723.960 43710.28 18677.400
## [4,]  7127.933 -2293.267 18677.40  8041.333

E = by(dat[,-5],dat[,5],function(x)(nrow(x)-1)*cov(x)) |> Reduce(f="+"); E
##         x1      x2      x3     x4
## x1 3895.62 1363.00 2462.46 564.50
## x2 1363.00 1696.20  812.08 480.84
## x3 2462.46  812.08 2722.26 627.18
## x4  564.50  480.84  627.18 615.66

Python尝试代码 #

import numpy as np
import pandas as pd
from sklearn import datasets
iris = datasets.load_iris() #iris dataset in sklearn for python
dat = pd.DataFrame(data=iris.data*10, columns=[f'x{i}' for i in range(1,5)] ) #same treatment as in R
dat["g"] = iris.target+1 #re-coding 3-level factor 
dat.shape
## (150, 5)
dat.head() #g=group, 3-level: 1,2,3
##      x1    x2    x3   x4  g
## 0  51.0  35.0  14.0  2.0  1
## 1  49.0  30.0  14.0  2.0  1
## 2  47.0  32.0  13.0  2.0  1
## 3  46.0  31.0  15.0  2.0  1
## 4  50.0  36.0  14.0  2.0  1

dat.groupby('g').size()
## g
## 1    50
## 2    50
## 3    50
## dtype: int64

dat.groupby('g').mean() #mean vector for data in each group; \bar{x}_i
##       x1     x2     x3     x4
## g
## 1  50.06  34.28  14.62   2.46
## 2  59.36  27.70  42.60  13.26
## 3  65.88  29.74  55.52  20.26

m = dat.iloc[:,:-1].mean(); m #mean vector for whole data; \bar{x}
## x1    58.433333
## x2    30.573333
## x3    37.580000
## x4    11.993333
## dtype: float64

# dat.groupby('g').cov()

H = dat.groupby('g').apply(lambda x: np.size(x)*(x.mean()-m)@(x.mean()-m).T).agg(sum,axis=0); H
E = dat.groupby('g').apply(lambda x: (np.size(x)-1)*np.cov(x)).agg(sum,axis=0); E

问题分析与修正代码 #

错误原因 #

1. 计算H时的问题：np.size(x)返回整个分组的元素总数（5列×50行=250），和R中nrow(x)取分组样本量（50）的逻辑不符；同时x.mean()-m是Series类型，转置后维度不匹配，需转为numpy数组做外积。
2. 计算E时的问题：同样误用np.size(x)获取样本量，应该用len(x)；且np.cov(x)默认按行计算协方差，需指定rowvar=False按列计算，和R的cov()行为对齐。

修正后的Python代码 #

import numpy as np
import pandas as pd
from sklearn import datasets

iris = datasets.load_iris()
dat = pd.DataFrame(data=iris.data*10, columns=[f'x{i}' for i in range(1,5)] )
dat["g"] = iris.target+1

# 转换为numpy数组方便后续矩阵运算
m = dat.iloc[:,:-1].mean().to_numpy()

# 计算组间平方和矩阵H
H = dat.groupby('g').apply(
    lambda x: len(x) * np.outer(x.iloc[:,:-1].mean().to_numpy() - m, 
                               x.iloc[:,:-1].mean().to_numpy() - m)
).sum()

# 计算组内平方和矩阵E
E = dat.groupby('g').apply(
    lambda x: (len(x)-1) * np.cov(x.iloc[:,:-1], rowvar=False)
).sum()

print("组间平方和矩阵H:\n", H)
print("\n组内平方和矩阵E:\n", E)

验证结果 #

运行后输出的H和E将与R语言的计算结果完全一致，数值无偏差。

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