df1:
交易日期 | 安全码 |
---|---|
20190319 | 000001 |
20220505 | 000001 |
df2:
付款日期 | 安全码 |
---|---|
20180708 | 000001 |
20190221 | 000002 |
df3:
结束日期 | 安全码 |
---|---|
20200203 | 000003 |
20210330 | 000004 |
我有3个数据框,结构如上所示。第一个数据帧的大小为 (3829931, 2),第二个 df 的大小为 (35689, 2),第三个 df 的大小为 (10329, 2)。我正在尝试使用 python 执行以下操作:
首先初始化一个名为“condition”的新列。对于 df1 的证券代码,如果在 df2 或 df3 的过去 2 年范围内可以找到对应的交易日期,则将条件设置为 1。例如,对于交易日期 20190319 的证券代码 000001,检查是否有df2 或 df3 中的一对,安全代码 = 000001,支付日期或结束日期在 20170319 和 20190319 之间。如果满足要求,则将条件设置为 1。
对于上面的示例数据帧,预期输出应该是第一行的条件为 1
我所做的尝试:
def check_condition(row):
# Calculate the date two years ago
two_years_ago = row['TRADEDATE'] - pd.DateOffset(years=2)
# Check if the code exists in df2 or df3 within the date range
condition_2 = ((df2['code'] == row['code']) &
(df2['PAYOUTDATE'] >= two_years_ago) &
(df2['PAYOUTDATE'] <= row['TRADEDATE'])).any()
condition_3 = ((df3['code'] == row['code']) &
(df3['ENDDATE'] >= two_years_ago) &
(df3['ENDDATE'] <= row['TRADEDATE'])).any()
# Return 1 if the condition is met, otherwise return 0
return 1 if condition_2 or condition_3 else 0
df1['condition'] = df1.apply(check_condition, axis=1)
以上是我写的代码。但是当我在本地环境中运行它时,速度非常慢。我想知道是否有更有效的方法来做到这一点?谢谢!
您可以使用向量化运算来代替循环,这样效率更高:
# Define the range of 2 years
timedelta = pd.to_datetime('today') - pd.DateOffset(years=2)
# Check conditions using vectorized operations
df2_conditions = ((df2['payoutDate'] >= timedelta) & (df2['payoutDate'] <= df1['tradeDate']))
df3_conditions = ((df3['endDate'] >= timedelta) & (df3['endDate'] <= df1['tradeDate']))
# Fill 'condition'
df1.loc[df2_conditions.any() | df3_conditions.any(), 'condition'] = 1
df1['condition'] = df1.loc[:, 'condition'].fillna(0)
# Drop empty rows
df1 = df1.dropna()
结果将是:
tradeDate securityCode condition
0 2019-03-19 000001 1.0
1 2022-05-05 000001 0.0