Pandas中的矢量化条件计数

Question

我有一个熊猫脚本，根据一些情况计算30天内再入院的次数。我想知道它是否可以矢量化以提高性能。我试过测向滚动（）。申请，但至今没有运气。你知道吗

下面是一个包含人工数据的表格来说明：

ID  VISIT_NO    ARRIVED      LEFT     HAD_A_MASSAGE   BROUGHT_A_FRIEND
1       1      29/02/1996  01/03/1996       0                 1
1       2      01/12/1996  04/12/1996       1                 0
2       1      20/09/1996  21/09/1996       1                 0
3       1      27/06/1996  28/06/1996       1                 0
3       2      04/07/1996  06/07/1996       0                 1
3       3      16/07/1996  18/07/1996       0                 1
4       1      21/02/1996  23/02/1996       0                 1
4       2      29/04/1996  30/04/1996       1                 0
4       3      02/05/1996  02/05/1996       0                 1
4       4      02/05/1996  03/05/1996       0                 1
5       1      03/10/1996  05/10/1996       1                 0
5       2      07/10/1996  08/10/1996       0                 1
5       3      10/10/1996  11/10/1996       0                 1

首先，我用ID创建一个字典：

ids = massage_df[massage_df['HAD_A_MASSAGE'] == 1]['ID']
id_dict = {id:0 for id in ids}

这张桌子上的每个人都做过按摩，但在我的真实数据集中，并不是所有人都这么幸运。你知道吗

接下来，我运行以下代码：

for grp, df in massage_df.groupby(['ID']):

    date_from = df.loc[df[df['HAD_A_MASSAGE']==1].index, 'LEFT']
    date_to = date_from + DateOffset(days=30)

    mask = ((date_from.values[0] < df['ARRIVED']) &
        (df['ARRIVED'] <= date_to.values[0]) &
        (df['BROGHT_A_FRIEND'] == 1))

    if len(df[mask]) > 0:
        id_dict[df['ID'].iloc[0]] = len(df[mask])

基本上，我想计算一个人最初来按摩的次数（单身或和朋友一起），然后在30天内和朋友一起回来。此表的预期结果将是IDs 3、IDs 4和IDs 5共6次再入院。你知道吗