在分组内按月份偏移DataFrame的value列（支持正负偏移）

问题背景 #

给定如下DataFrame：

import pandas as pd

df = pd.DataFrame({
    "id": [0, 0, 0, 1, 1, 1, 1, 2, 2, 2],
    "date": ["2017-01-31", "2017-02-28", "2017-03-31", "2017-01-31", "2017-03-31", "2017-04-30", "2017-05-31", "2017-01-31", "2017-03-31", "2017-05-31"],
    "value": [10., 12., 15., 8., 11., 15., 17., 6., 14., 15.]
})

df["date"] = pd.to_datetime(df["date"], format="%Y-%m-%d")

需要在每个id分组内，根据指定偏移量列表（支持正负值）生成value的月度偏移列，比如偏移量[1,2]对应生成value_1和value_2列，期望结果对应的表格如下：

id  date        value   value_1  value_2
0   0   2017-01-31  10.0    12.0     15.0
1   0   2017-02-28  12.0    15.0     NaN
2   0   2017-03-31  15.0    NaN      NaN
3   1   2017-01-31  8.0     NaN      11.0
4   1   2017-03-31  11.0    15.0     17.0
5   1   2017-04-30  15.0    17.0     NaN
6   1   2017-05-31  17.0    NaN      NaN
7   2   2017-01-31  6.0     NaN      14.0
8   2   2017-03-31  14.0    NaN      15.0
9   2   2017-05-31  15.0    NaN      NaN

当前实现方式可行，但数据量较大且偏移列表包含多个元素时效率偏低：

from pandas.tseries.offsets import MonthEnd

shifts = [1, 2]

tmp = df.copy()
for shift in shifts:
    tmp_shifed = tmp.rename(columns={"value": f"value_{shift}"}).assign(date=df["date"] + MonthEnd(-1 * shift))
    df = df.merge(tmp_shifed, on=["id", "date"], how="left")

高效解决方案 #

方法1：分组时间序列对齐+批量Shift #

核心思路是将每个分组转换为连续月度的时间序列，通过批量Shift操作生成所有偏移列，避免多次循环Merge：

import pandas as pd

shifts = [1, 2]

# 1. 按id分组，将日期转为月度频率索引，确保时间序列连续
grouped = df.groupby('id').apply(
    lambda x: x.set_index('date')['value'].asfreq('M')
).unstack(0)

# 2. 批量生成偏移列，shift(-s)对应取未来s个月的value
shifted_dfs = []
for s in shifts:
    shifted = grouped.shift(-s).stack().reset_index(name=f'value_{s}')
    shifted_dfs.append(shifted)

# 3. 合并所有偏移列回原表
result = df.copy()
for shifted_df in shifted_dfs:
    result = result.merge(shifted_df, on=['id', 'date'], how='left')

print(result)

方法2：向量化一次性处理所有偏移 #

通过生成所有偏移日期的映射表，仅用一次Merge完成所有匹配，减少循环开销：

import pandas as pd
from pandas.tseries.offsets import MonthEnd

shifts = [1, 2]

# 1. 生成包含所有偏移日期的临时表
tmp = df.assign(
    **{f'shift_{s}': df['date'] + MonthEnd(s) for s in shifts}
).melt(
    id_vars=['id', 'date', 'value'],
    value_vars=[f'shift_{s}' for s in shifts],
    var_name='shift',
    value_name='target_date'
).assign(
    shift=lambda x: x['shift'].str.extract(r'shift_(\d+)').astype(int)
)

# 2. 匹配id和目标日期，获取对应value
merged = tmp.merge(
    df[['id', 'date', 'value']],
    left_on=['id', 'target_date'],
    right_on=['id', 'date'],
    how='left',
    suffixes=('', '_shifted')
)

# 3. 将偏移列转换为宽表格式
result = merged.pivot(
    index=['id', 'date', 'value'],
    columns='shift',
    values='value_shifted'
).rename(columns=lambda x: f'value_{x}').reset_index()

print(result)

方案优势 #

• 方法1利用pandas分组时间序列特性，避免多次IO密集型的Merge操作，大数据量下效率提升明显；
• 方法2通过向量化操作一次性处理所有偏移逻辑，循环次数降至最低；
• 两种方案均支持正负偏移量，只需调整Shift方向（正偏移用shift(s)，负偏移用shift(-s)）。

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