如何用Pandas按组计算条件行前后每n天的斜率

问题背景 #

给定带分组的Pandas DataFrame，需按group字段分组，以每组中indicator == True的行作为基准行，完成以下操作：

• 给基准行前后的行按每n行划分为分组，生成id字段（基准行id=0，基准前分组为负整数，基准后为正整数）
• 计算每个分组内diff_days（x轴）与value（y轴）的线性回归斜率，赋值给组内所有行的slope字段（基准行slope=0）

原始数据代码 #

import pandas as pd

n = 3

data = [['A', '2022-09-01', False, 2, -3], ['A', '2022-09-02', False, 1, -2], ['A', '2022-09-03', False, 1, -1], ['A', '2022-09-04', True, 3, 0], 
        ['A', '2022-09-05', False, 3, 1], ['A', '2022-09-06', False, 2, 2], ['A', '2022-09-07', False, 1, 3], ['A', '2022-09-07', False, 2, 3], 
        ['A', '2022-09-08', False, 4, 4], ['A', '2022-09-09', False, 2, 5],
        ['B', '2022-09-01', False, 2, -4], ['B', '2022-09-02', False, 2, -3], ['B', '2022-09-03', False, 4, -2], ['B', '2022-09-04', False, 2, -1], 
        ['B', '2022-09-05', True, 2, 0], ['B', '2022-09-06', False, 2, 1], ['B', '2022-09-07', False, 1, 2], ['B', '2022-09-08', False, 3, 3], 
        ['B', '2022-09-09', False, 3, 4], ['B', '2022-09-10', False, 2, 5]]
df = pd.DataFrame(data=data, columns=['group', 'date', 'indicator', 'value', 'diff_days'])

需求示例说明（以分组A为例） #

• 基准行（行3）前的3行（行0-2）为id=-1组，斜率为slope(x=[-3,-2,-1], y=[2,1,1])=-0.5
• 基准行后的前3行（行4-6）为id=1组，斜率为slope(x=[1,2,3], y=[3,2,1])=-1
• 基准行后的剩余3行（行7-9）为id=2组，斜率为slope(x=[3,4,5], y=[2,4,2])=0

期望输出代码 #

data = [['A', '2022-09-01', False, 2, -3, -1, -0.5], ['A', '2022-09-02', False, 1, -2, -1, -0.5], ['A', '2022-09-03', False, 1, -1, -1, -0.5], ['A', '2022-09-04', True, 3, 0, 0, 0], 
        ['A', '2022-09-05', False, 3, 1, 1, -1], ['A', '2022-09-06', False, 2, 2, 1, -1], ['A', '2022-09-07', False, 1, 3, 1, -1], ['A', '2022-09-07', False, 2, 3, 2, 0], 
        ['A', '2022-09-08', False, 4, 4, 2, 0], ['A', '2022-09-09', False, 2, 5, 2, 0],
        ['B', '2022-09-01', False, 2, -4, -2, 0], ['B', '2022-09-02', False, 2, -3, -1, 0], ['B', '2022-09-03', False, 4, -2, -1, 0], ['B', '2022-09-04', False, 2, -1, -1, 0], 
        ['B', '2022-09-05', True, 2, 0, 0, 0], ['B', '2022-09-06', False, 2, 1, 1, 0.5], ['B', '2022-09-07', False, 1, 2, 1, 0.5], ['B', '2022-09-08', False, 3, 3, 1, 0.5], 
        ['B', '2022-09-09', False, 3, 4, 2, -1], ['B', '2022-09-10', False, 2, 5, 2, -1]]
df_desired = pd.DataFrame(data=data, columns=['group', 'date', 'indicator', 'value', 'diff_days', 'id', 'slope'])

实现代码 #

import pandas as pd
import numpy as np

# 原始数据（可替换为自定义数据集）
n = 3
data = [['A', '2022-09-01', False, 2, -3], ['A', '2022-09-02', False, 1, -2], ['A', '2022-09-03', False, 1, -1], ['A', '2022-09-04', True, 3, 0], 
        ['A', '2022-09-05', False, 3, 1], ['A', '2022-09-06', False, 2, 2], ['A', '2022-09-07', False, 1, 3], ['A', '2022-09-07', False, 2, 3], 
        ['A', '2022-09-08', False, 4, 4], ['A', '2022-09-09', False, 2, 5],
        ['B', '2022-09-01', False, 2, -4], ['B', '2022-09-02', False, 2, -3], ['B', '2022-09-03', False, 4, -2], ['B', '2022-09-04', False, 2, -1], 
        ['B', '2022-09-05', True, 2, 0], ['B', '2022-09-06', False, 2, 1], ['B', '2022-09-07', False, 1, 2], ['B', '2022-09-08', False, 3, 3], 
        ['B', '2022-09-09', False, 3, 4], ['B', '2022-09-10', False, 2, 5]]
df = pd.DataFrame(data=data, columns=['group', 'date', 'indicator', 'value', 'diff_days'])

# 计算线性回归斜率的工具函数
def get_slope(x_values, y_values):
    if len(x_values) <= 1:
        return 0.0
    x_arr = np.array(x_values)
    y_arr = np.array(y_values)
    slope, _ = np.polyfit(x_arr, y_arr, 1)
    return round(slope, 1)

# 按分组处理逻辑
processed_dfs = []
for group_name, sub_df in df.groupby('group'):
    # 定位基准行
    base_row_idx = sub_df[sub_df['indicator']].index[0]
    base_row = sub_df.loc[base_row_idx].copy()
    base_row['id'] = 0
    base_row['slope'] = 0.0
    
    # 拆分基准前后的数据
    before_base = sub_df.loc[:base_row_idx-1].copy()
    after_base = sub_df.loc[base_row_idx+1:].copy()
    
    # 处理基准前的分组：从靠近基准的开始，每n行一组，id为负
    if not before_base.empty:
        reversed_before = before_base.iloc[::-1]
        reversed_before['id'] = -(np.arange(len(reversed_before)) // n + 1)
        before_base = reversed_before.iloc[::-1]
    
    # 处理基准后的分组：从靠近基准的开始，每n行一组，id为正
    if not after_base.empty:
        after_base['id'] = np.arange(len(after_base)) // n + 1
    
    # 合并当前分组的所有行
    combined = pd.concat([before_base, base_row.to_frame().T, after_base], ignore_index=True)
    
    # 计算每个(group, id)组的斜率，并映射回每行
    slope_groups = combined.groupby(['group', 'id']).apply(
        lambda g: get_slope(g['diff_days'], g['value'])
    ).reset_index(name='slope')
    
    combined = combined.merge(slope_groups, on=['group', 'id'], how='left')
    processed_dfs.append(combined)

# 合并所有分组的结果，调整列顺序
final_df = pd.concat(processed_dfs, ignore_index=True)
final_df = final_df[['group', 'date', 'indicator', 'value', 'diff_days', 'id', 'slope']]

# 验证结果（可选）
print(final_df)

关键逻辑说明 #

1. 斜率计算：使用np.polyfit计算线性回归斜率，处理了分组内仅一行的边界情况（直接返回0）
2. 基准前分组：通过反转数据，从靠近基准的行开始分组，保证id=-1是最接近基准的前n行
3. 分组id生成：利用整数除法//快速生成连续的分组标识，高效简洁
4. 斜率映射：先按group和id分组计算斜率，再通过合并将斜率赋值给组内所有行，避免重复计算

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