创建基于其他数据帧行应用标志的函数

Question

我有一个数据帧，看起来像这个

date            id      type
02/02/2020      2       A
29/02/2020      2       B
04/03/2020      2       B
02/01/2020      3       B 
15/01/2020      3       A 
19/01/2020      3       C
...             ...     ...

我想创建一个新列，称为flagged。对于每一行，如果存在具有的另一行，我希望flagged的值等于True

• 相同的id
• A类
• 与当前行的日期相差大于0且小于30的日期

我想把上面的数据帧转换成这个

date            id      type    flagged
02/02/2020      2       A       False
29/02/2020      2       B       True
04/03/2020      2       B       False
02/01/2020      3       B       False
15/01/2020      3       A       False
19/01/2020      3       C       True
...             ...     ...     ...

我的方法：

我创建了以下功能

def check_type(id, date):
if df[(df.id == id) & (df.type == 'A') & (date - df.date > datetime.timedelta(0)) & (date - df.date < datetime.timedelta(30))].empty:
return False
else:
return True

这样，如果我运行

df['flagged'] = df.apply(lambda x: check_type(x.id, x.date), axis = 1)

我得到了想要的结果。

问题：

1. 如何更改函数check_type，使其适用于任何数据帧，无论其名称如何？当前函数只有在其使用的数据帧被称为df时才起作用
2. 如何加快此过程？我想在一个大的数据帧上运行这个函数，但它的执行速度没有我想要的那么快

提前感谢！

Answer 1

我会用A类型找到最后一个日期，并用ffill在整个id中传播它，然后找到差异：

last_dates = df.date.where(df['type'].eq('A')).groupby(df['id']).ffill()
# this is the new column
df.date.sub(last_dates).lt(pd.to_timedelta('30D')) & df['type'].ne('A')

输出：

0    False
1     True
2    False
3    False
4    False
5     True
dtype: bool

注意：这是有效的，因为总是用False屏蔽A。