DataFrame分组加权平均：行数不同时如何正确计算分母权重和？

先看原始的DataFrame和初始代码：

import pandas as pd
import numpy as np

data = [
  ['A',1,2,3,4],
  ['A',5,6,7,8],
  ['A',9,10,11,12],
  ['B',13,14,15,16],
  ['B',17,18,19,20],
  ['B',21,22,23,24],
  ['B',25,26,27,28],
  ['C',29,30,31,32],
  ['C',33,34,35,36],
  ['C',37,38,39,40],
  ['D',13,14,15,0],
  ['D',0,18,19,0],
  ['D',0,0,23,0],
  ['D',0,0,0,0],
  ['E',13,14,15,0],
  ['E',0,18,19,0],
  ['F',0,0,23,0],
]

df = pd.DataFrame(data, columns=['Name', 'num1', 'num2', 'num3', 'num4'])

初始计算分组加权平均的代码：

weights = [10,20,30,40] 

df=df.groupby('Name').agg(lambda g: sum(g*weights[:len(g)])/sum(weights[:len(g)]))

问题分析 #

当前代码中weights[:len(g)]是取权重列表的前N个元素（N为分组的行数），但实际需求是取后N个元素：

• A组3行，当前用[10,20,30]，预期要用[20,30,40]
• E组2行，当前用[10,20]，预期要用[30,40]

修正后的代码 #

只需要把切片从weights[:len(g)]改成weights[-len(g):]，负索引会从列表末尾开始取对应数量的元素：

weights = [10,20,30,40] 

df = df.groupby('Name').agg(lambda g: sum(g * weights[-len(g):]) / sum(weights[-len(g):]))

验证示例 #

• A组num1的计算：(1*20 +5*30 +9*40)/(20+30+40) = (20+150+360)/90 = 530/90 ≈5.888，符合预期
• E组num1的计算：(13*30 +0*40)/(30+40) = 390/70≈5.571，符合预期

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