如何通过pandas groupby分组计算为DataFrame生成符合规则的新列

问题原因 #

你原有代码的报错和异常来自两个问题：

1. 自定义函数中for info in group实际是遍历分组的列名而非行，循环次数和分组行数不匹配，导致后续索引取值报错
2. 不需要调用join拼接结果，groupby.apply直接会返回和原表索引对齐的完整DataFrame，额外拼接会出现列重复问题

最优实现方案（矢量化操作，效率远高于循环） #

直接用pandas内置的shift方法取分组内前一行的Red值，计算逻辑一行就能完成：

import pandas as pd

df = pd.DataFrame({
    'Red' : [1,2,3,4,5,6,7,8,9,10],
    'Groups':['A','B','A','A','B','C','B','C','B','C'],
    'Blue':[10,20,30,40,50,60,70,80,90,100]
})

# 按Groups分组后，将Red列下移一行，取每个分组前一行的Red值
df['Total'] = df['Blue'] - df.groupby('Groups')['Red'].shift(1)
# 分组首行的空值填0，转换为整数格式
df['Total'] = df['Total'].fillna(0).astype(int)

运行后得到的Total列和你给出的预期结果完全一致。

自定义函数修改方案 #

如果要保留你原有循环实现的写法，修改如下即可正常运行：

def funct(group):
    lst = []
    # 直接遍历分组的行索引，而不是遍历列名
    for count in range(len(group)):
        if count == 0:
            lst.append(0)
        else: 
            num = group.iloc[count]['Blue'] - group.iloc[count-1]['Red']
            lst.append(num)
    group['Total'] = lst
    return group

# 加group_keys=False避免生成多余的分组索引，不需要join直接接收返回结果
df = df.groupby('Groups', group_keys=False).apply(funct)

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