Pandas如何按学生分组计算成绩斜率并新增为列

你现有代码的问题是没有按学生维度分组，直接对全量数据做线性回归，得到的是全局斜率，而非每个学生各自的成绩随时间变化的斜率。
最高效简洁的实现方式是使用pandas原生的groupby分组计算逻辑，无需手动循环或额外合并表，代码如下：

import pandas as pd
import scipy.stats
import datetime as dt

# 构造原始数据集
df = pd.DataFrame({'student': 'A A A B B B B'.split(),
                  'exam_date': pd.date_range(start='1/1/2020', periods=7, freq='M'),
                  'score': [15, 28, 17, 22, 43, 40, 52]})

# 生成日期对应的序数值，作为线性回归的数值型自变量
df['date_ordinal'] = df['exam_date'].map(dt.datetime.toordinal)

# 按学生分组计算回归斜率，直接广播到原表生成新列
df['slope'] = df.groupby('student')\
                .apply(lambda g: scipy.stats.linregress(g['date_ordinal'], g['score']).slope)\
                .round(3)\
                .loc[df['student']]\
                .values

# 不需要中间辅助列的话可以直接删掉
# df = df.drop(columns='date_ordinal')

运行后输出结果和你预期完全一致：

student  exam_date  score  slope
0       A 2020-01-31     15  0.028
1       A 2020-02-29     28  0.028
2       A 2020-03-31     17  0.028
3       B 2020-04-30     22  0.285
4       B 2020-05-31     43  0.285
5       B 2020-06-30     40  0.285
6       B 2020-07-31     52  0.285

方案说明 #

• 不需要重复调用pd.to_datetime转换exam_date字段，通过pd.date_range生成的列本身就是datetime类型，重复转换属于冗余操作
• 分组计算后通过.loc[df['student']]直接将每个学生对应的斜率匹配回原表所有行，比merge操作性能更高，代码更短
• 链式调用.round(3)直接将斜率保留3位小数，和预期输出格式对齐
• 如果不需要保留用于计算的date_ordinal辅助列，计算完成后直接删除即可

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