Pandas查找同ID同状态前序行 按日期差计算行权重

现有数据 #

当前持有如下结构的DataFrame：

id      date        status
--------------------------------
0    2     2019-02-04   reserve    
1    2     2020-01-01   reserve    
2    2     2020-01-02   sold      
3    3     2020-01-03   reserve    
4    4     2020-01-03   booked     
5    3     2020-02-05   reserve    
6    4     2020-02-06   sold       
7    3     2020-02-07   sold

计算规则 #

需要为每一行新增weight列，规则如下：

• 若当前行满足status == 'sold'，weight直接赋值为1
• 其余状态的行：查找和当前行id相同、status相同的最近前序行，取两行日期的天数差，按1 / 天数差计算权重；无匹配前序行时权重为0

期望输出 #

最终得到的DataFrame效果如下：

id      date        status    weight
----------------------------------------
0    2     2019-02-04   reserve     0 
1    2     2020-01-01   reserve   0.003 
2    2     2020-01-02   sold        1
3    3     2020-01-03   reserve     0 
4    4     2020-01-03   booked      0 
5    3     2020-02-05   reserve   0.030  
6    4     2020-02-06   sold        1 
7    3     2020-02-07   sold        1

现有代码问题 #

当前编写的代码无法正确实现带条件的前序行查找，代码如下：

df.date = pd.to_datetime(df.date)
df.sort_values(['date', 'id'], inplace=True)
df = df.reset_index(drop='index')

df['weight'] = np.where(df['status'] == 'sold', 1, 0)
df = df.reset_index(drop='index')
df1 = df[~(df.status == 'sold')]
retain_index = list(df1.index.values)
df1 = df1.reset_index(drop='index')

df1['diff_in_dates'] = df1.groupby('flat_id')['date'].diff().dt.days.fillna(0)

for index, row in df1.iterrows():
    if index == 0:
        pass
    else:
        if df1.loc[index, 'weight'] == 0:
            if (df1.loc[index, 'status'] == df1.loc[index - 1, 'status']) & \
                                    (df1.loc[index, 'flat_id'] == df1.loc[index - 1, 'flat_id']):
                df1.loc[index, 'weight'] = 1 / df1.loc[index, 'diff_in_dates']

df1.weight.replace([np.inf, -np.inf], np.nan, inplace=True)
df1.weight = df1.weight.fillna(0)

# 此处原df索引为5的行计算得到weight=0，不符合预期的0.030
df = pd.DataFrame(df.append(df1, ignore_index = False)

代码计算错误点：原表索引为5的行权重计算结果为0，不符合预期的0.030，需要修正逻辑满足需求。

修正实现 #

原代码存在三个核心问题：

• 分组字段笔误，将id错写为不存在的flat_id
• 分组维度缺失，仅按id分组没有叠加status条件，无法匹配同id同状态的前序行
• 仅对比相邻行、拆分df再拼接的逻辑冗余，容易出现索引匹配错误

可以直接用pandas原生groupby方法实现，无需循环：

import pandas as pd
import numpy as np

# 预处理日期格式
df['date'] = pd.to_datetime(df['date'])

# 先给sold状态行赋权重1
df['weight'] = np.where(df['status'] == 'sold', 1, 0)

# 筛选非sold行，按id+status双维度分组计算日期间隔
non_sold_flag = df['status'] != 'sold'
day_diff = df[non_sold_flag].groupby(['id', 'status'])['date'].diff().dt.days

# 计算非sold行权重，无前序行的空值填0，保留3位小数和预期格式对齐
df.loc[non_sold_flag, 'weight'] = (1 / day_diff).fillna(0).round(3)

# 按日期、id排序恢复原表顺序
df = df.sort_values(['date', 'id']).reset_index(drop=True)

运行后结果和预期完全一致：id=3、status=reserve的2020-02-05行和同组前序行2020-01-03间隔33天，1/33≈0.030，计算正确；id=2的两条reserve行间隔331天，1/331≈0.003，符合预期。

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