如何在Pandas中为每个ID统一填充参考日期后的最小时间戳

需求与问题 #

我有一个Pandas DataFrame，想要创建min_date_after_ref_date列，展示每个ID在指定参考日期ref_date之后的最小时间戳。

原代码如下：

from datetime import datetime
import pandas as pd

ref_date = datetime.strptime('2023-04-21 12:00', '%Y-%m-%d %H:%M')
df = pd.DataFrame({'id':[1,2,1,1,3], 'time_stamp': ['2023-04-19 12:05', '2023-04-21 12:45',
                                                 '2023-04-21 15:45', '2023-04-23 13:15', '2023-04-18 12:05']})
df = df.assign(time_stamp=pd.to_datetime(df.time_stamp))
df = df.assign(min_date_after_ref_date=df[df.time_stamp>ref_date].groupby('id').time_stamp.transform('min'))

运行后得到结果：

id           time_stamp min_date_after_ref_date
0   1 2023-04-19 12:05:00                      NaT
1   2 2023-04-21 12:45:00        2023-04-21 12:45:00
2   1 2023-04-21 15:45:00        2023-04-21 15:45:00
3   1 2023-04-23 13:15:00        2023-04-21 15:45:00
4   3 2023-04-18 12:05:00                      NaT

现在希望每个ID对应的min_date_after_ref_date值保持一致，比如第一行（id=1）也显示2023-04-21 15:45:00而非NaT，该如何修改代码？

解决方案 #

原代码的问题在于：先筛选出time_stamp>ref_date的行再分组，导致原DataFrame中不满足条件的行无法匹配到分组后的结果，只能显示NaT。

以下两种方法可以解决这个问题：

方法1：先计算分组结果再映射 #

先按id分组计算每个组内满足条件的最小时间戳，再通过map将结果批量映射回原DataFrame：

from datetime import datetime
import pandas as pd

ref_date = datetime.strptime('2023-04-21 12:00', '%Y-%m-%d %H:%M')
df = pd.DataFrame({'id':[1,2,1,1,3], 'time_stamp': ['2023-04-19 12:05', '2023-04-21 12:45',
                                                 '2023-04-21 15:45', '2023-04-23 13:15', '2023-04-18 12:05']})
df['time_stamp'] = pd.to_datetime(df['time_stamp'])

# 预计算每个id的目标最小日期
min_dates = df[df['time_stamp'] > ref_date].groupby('id')['time_stamp'].min()
# 将结果映射到原DataFrame的对应id行
df['min_date_after_ref_date'] = df['id'].map(min_dates)

方法2：分组内直接处理 #

利用groupby+transform，在每个分组内部先筛选符合条件的时间戳再取最小值，确保同一id的所有行都能拿到结果：

from datetime import datetime
import pandas as pd

ref_date = datetime.strptime('2023-04-21 12:00', '%Y-%m-%d %H:%M')
df = pd.DataFrame({'id':[1,2,1,1,3], 'time_stamp': ['2023-04-19 12:05', '2023-04-21 12:45',
                                                 '2023-04-21 15:45', '2023-04-23 13:15', '2023-04-18 12:05']})
df['time_stamp'] = pd.to_datetime(df['time_stamp'])

df['min_date_after_ref_date'] = df.groupby('id')['time_stamp'].transform(
    lambda x: x[x > ref_date].min() if (x > ref_date).any() else pd.NaT
)

修改后的结果：

id           time_stamp min_date_after_ref_date
0   1 2023-04-19 12:05:00        2023-04-21 15:45:00
1   2 2023-04-21 12:45:00        2023-04-21 12:45:00
2   1 2023-04-21 15:45:00        2023-04-21 15:45:00
3   1 2023-04-23 13:15:00        2023-04-21 15:45:00
4   3 2023-04-18 12:05:00                      NaT

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