在Pandas中矢量化条件计数

Question

我有一个Pandas脚本，根据一些条件计算30天内入院的次数。我想知道它是否可以进行矢量化以提高性能。我已经尝试过df.rolling（）。apply，但到目前为止没有运气。

这是一张包含人为数据的表格，用于说明：

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天内和朋友一起回来。该表的预期结果将是ID 3,4和5的总共6次重新接收。