python - 根据后续几行的值检查数据帧中每一行的条件 - 代码非常慢

Question

在包含列 person_id (int)、dates (datetime) 和 is_hosp 的数据帧 df 中( bool 值)，我需要找出每个 person_id 的每个日期，在该行日期之后的接下来 15 天内，is_hosp 是否为 true。

编辑:

根据对原始问题的相关评论，我应该指定每行/person_id 的观察日期并不总是连续的(即 person_id 可以对 01/01/2018 进行观察，然后对2018 年 8 月 1 日)。

我写了下面的代码，但是速度很慢。

为了重现该问题，这里有一个与我正在处理的数据帧格式相同的假数据帧:

import pandas as pd
import datetime as dt
import numpy as np

count_unique_person = 10
dates = pd.date_range(dt.datetime(2019, 1, 1), dt.datetime(2021, 1, 1))
dates = dates.to_pydatetime().tolist()
person_id = [i for i in range(1, count_unique_person + 1) for _ in range(len(dates))]
dates = dates * count_unique_person
sample_arr = [True, False]
col_to_check = np.random.choice(sample_arr, size = len(dates))
df = pd.DataFrame()
df['person_id'] = person_id
df['dates'] = dates
df['is_hosp'] = col_to_check

以下是我实现检查的方法(is_inh_15d 列)，但运行时间太长(我的原始数据帧包含一百万行):

is_inh = []
is_inh_idx = []
for each_p in np.unique(df['person_id'].values):
    t_df = df[df['person_id'] == each_p].copy()
    for index, e_row in t_df.iterrows():
        s_dt = e_row['dates']
        e_dt = s_dt + dt.timedelta(days = 15)
        t_df2 = t_df[(t_df['dates'] >= s_dt) & (t_df['dates'] <= e_dt)]
        is_inh.append(np.any(t_df2['is_hosp'] == True))
        is_inh_idx.append(index)
h_15d_df = pd.DataFrame()
h_15d_df['is_inh_15d'] = is_inh
h_15d_df['is_inh_idx'] = is_inh_idx
h_15d_df.set_index('is_inh_idx', inplace = True)
df = df.join(h_15d_df)

我不知道如何矢量化检查每行接下来 15 天的逻辑，以查看“is_hosp”是否为 True。

有人可以建议一下吗？

Answer 1

.rolling() 在时间序列数据或类似日期时间的索引上接受 16 天的 "16D"。 (16天是指连续15天及其后的15天)

is_inh_15d = (
    df.sort_values('dates', ascending=False)
    .groupby('person_id')
    .rolling('16D', on='dates')
    .is_hosp.max() == 1
)
df2 = df.join(is_inh_15d, on=['person_id', 'dates'], rsuffix='_15d')