pandas基于日期差与序列条件非循环分配观测值权重

• 此前就同一问题在Stack Overflow提问时，用户@mozway提供了大量帮助，但之前使用的权重分配逻辑存在错误。
• 需求是为DataFrame生成weight列，目标数据格式如下：

id      date        status    weight   diff_in_days  comment
-----------------------------------------------------------------
0    2     2019-02-03   reserve   0.003         0       1 / diff_days
1    2     2019-12-31   reserve   0.001       331       since diff to next is 1 day 
2    2     2020-01-01   reserve    0.9          1       since the next date status is sold   
3    2     2020-01-02   sold        1           1       sold
4    3     2020-01-03   reserve   0.001         0       since diff to next is 1 day
5    4     2020-01-03   booked     0.9          0       since the next date status is sold
6    3     2020-02-04   reserve    0.9          1       since the next date status is sold
7    4     2020-02-06   sold        1           3       sold
7    3     2020-02-07   sold        1           3       sold

• 目前使用以下代码生成diff_in_days列：

df['diff_in_days'] = df.groupby('flat_id')['date'].diff().dt.days.fillna(0)

需要找到无需for循环的实现方案，对应逻辑的伪代码如下：

for i in df.iterrows():
    df['weight'][i] = 1 / df['diff_in_days'][i+1]
if df['status'][i+1] == 'sold' (for each flat_id):
    df['weight'][i] = 0.9
if df['status'][i] == 'sold':
    df['weight'][i] = 1

直接使用pandas向量化操作即可实现，不需要写循环，步骤如下：

1. 按id分组，通过shift(-1)把同组下一行的日期间隔、下一行的状态匹配到当前行
2. 按规则优先级赋值，优先级从高到低：
   • 当前行状态为sold时，权重固定为1
   • 同组下一行状态为sold时，权重固定为0.9
   • 其余情况，权重为1/当前行与下一行的日期间隔，注意提前把间隔为0的情况做替换，避免除0错误，可根据业务规则调整兜底值

对应可直接运行的代码：

import numpy as np

# 分组提取下一行的辅助字段
grouped = df.groupby('id')
df['next_day_diff'] = grouped['diff_in_days'].shift(-1)
df['next_status'] = grouped['status'].shift(-1)

# 按优先级批量赋值
df['weight'] = np.where(
    df['status'] == 'sold',
    1,
    np.where(
        df['next_status'] == 'sold',
        0.9,
        1 / df['next_day_diff'].replace(0, 333) # 此处0替换为333时计算结果约为0.003，和示例首行值匹配，可按需调整
    )
)

# 删除临时生成的辅助列
df.drop(columns=['next_day_diff', 'next_status'], inplace=True)

逻辑匹配说明：当前生成的diff_in_days是同组内当前行和上一行的日期差，因此下一行的diff_in_days值，恰好就是当前行和下一行的日期间隔，和伪代码中取i+1位置diff_in_days的逻辑完全一致。

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