如何优雅实现Pandas分组均值DataFrame与字典按键匹配相乘

问题复现 #

基础数据构造与分组均值计算代码如下：

import pandas as pd
import numpy as np

TownA = [35, 43, 36, 39, 28, 28, 29, 25, 38, 27, 26, 32, 29, 40, 35, 41, 37, 31, 45, 34]
TownB = [27, 15, 4, 41, 49, 25, 10, 30]
Rural = [8, 14, 12, 15, 30, 32, 21, 20, 34, 7, 11, 24]
tv_viewing = TownA + TownB + Rural

labels = np.array(["TownA", "TownB", "Rural"])
stratum = np.repeat(labels, [len(TownA), len(TownB), len(Rural)], axis=0)
sample = pd.DataFrame()
sample["viewership"] = tv_viewing
sample["stratum"] = stratum

grouped_mean = sample.groupby("stratum").mean()

分组均值计算结果：

待匹配的分层人口字典：

Population = {"TownA": 155, "TownB": 62, "Rural": 93}

需求为按分层键匹配字典值，和对应分组的收视均值相乘，得到如下加权结果：

原有实现通过重置索引、字典转DataFrame、merge合并后计算乘积，流程冗余。尝试编写的apply代码触发'Series' object has no attribute 'viewership'报错。

报错原因 #

DataFrame.apply()默认按列遍历，传入lambda函数的参数是整列的Series对象，不是单行数据，因此不存在viewership、stratum属性。且分组后stratum已经是DataFrame的索引，不在列中，即使按行遍历也无法通过x.stratum取值。

简洁实现方案 #

方案1：索引映射（最优，性能最好） #

不需要重置索引、不需要合并表，直接通过索引匹配字典值做向量运算即可：

grouped_mean = sample.groupby("stratum").mean()
# 匹配索引对应的人口数，和收视均值相乘得到结果，赋值为新列
grouped_mean["weighted_result"] = grouped_mean["viewership"] * grouped_mean.index.map(Population)

方案2：按行apply实现 #

如果要使用apply，需要指定axis=1参数切换为按行遍历，行的索引值通过name属性获取：

grouped_mean = sample.groupby("stratum").mean()
grouped_mean["weighted_result"] = grouped_mean.apply(
    lambda row: row["viewership"] * Population[row.name], 
    axis=1
)

两种方案运行后都可以直接得到期望的计算结果，无多余操作。

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