Python:如何最有效地向量化这个forloop(使用pandas groupby)

2024-06-25 23:06:18 发布

您现在位置:Python中文网/ 问答频道 /正文

我正在尝试创建一个聚合股票函数,在日期上循环。基本上,我希望基于分组和迭代日期对变量求和。但是,我在给定的一天中有多个条目,这使问题稍微复杂化了。我目前正在使用pandas.groupby,但这对我的>;来说太慢了;3m行,所以我想知道如何将其矢量化

由一些示例代码表示,我想将以下内容矢量化。假设我正在为一组客户储存一些东西,比如苹果,我想记录下在任何一天我为每个客户储存了多少苹果。每个客户可以从一天到将来的某个日期“存款”一些苹果,并且每天可以存款几笔。 示例代码如下所示:

import pandas as pd

df = {'StartDate': ['2020-01-01','2020-01-01','2020-01-01','2020-01-02','2020-01-02','2020-01-02','2020-01-03','2020-01-04','2020-01-04','2020-01-05'],
      'EndDate':['2020-01-02','2020-01-02','2020-01-05','2020-01-05','2020-01-03','2020-01-04','2020-01-06','2020-01-06','2020-01-06','2020-01-06'],
      'Apples':[5,6,2,4,4,10,8,9,3,7],
      'Customer':['A','B','A','C','A','A','B','B','C','A']}
df = pd.DataFrame(data=df)

cycledates = ['2020-01-01','2020-01-02','2020-01-03','2020-01-04','2020-01-05','2020-01-06']

def stockfunction(indata, ingroupingvars, cycledates):
    stock_agg = pd.DataFrame()
    for i in cycledates:
        stock_input = indata.loc[
            (indata['StartDate'] <= i) & (indata['EndDate'] > i)]
        stock_input['EvaluationDate'] = i
        stock_marg = stock_input.groupby(ingroupingvars)['Apples'].sum().unstack(level=0)
        stock_agg = stock_agg.append(stock_marg)
    return stock_agg

ingroupingvars = ['Customer','EvaluationDate']
stock = stockfunction(df,ingroupingvars,cycledates)

样本输出

数据帧如下所示:

StartDate     EndDate  Apples Customer
0  2020-01-01  2020-01-02       5        A
1  2020-01-01  2020-01-02       6        B
2  2020-01-01  2020-01-05       2        A
3  2020-01-02  2020-01-05       4        C
4  2020-01-02  2020-01-03       4        A
5  2020-01-02  2020-01-04      10        A
6  2020-01-03  2020-01-06       8        B
7  2020-01-04  2020-01-06       9        B
8  2020-01-04  2020-01-06       3        C
9  2020-01-05  2020-01-06       7        A

最终结果如下所示:

Customer         A     B    C
EvaluationDate               
2020-01-01       7   6.0  NaN
2020-01-02      16   NaN  4.0
2020-01-03      12   8.0  4.0
2020-01-04       2  17.0  7.0
2020-01-05       7  17.0  3.0

…只是在我的原始数据上运行要花很长时间


Tags: 苹果dfinput客户stockcustomeraggpd
1条回答
网友
1楼 · 发布于 2024-06-25 23:06:18

看起来您要执行交叉合并和查询:

# convert to datetime type for safe comparison
Eval_df = pd.to_datetime(pd.Series(cycledates, name='EvaluationDate')).to_frame()

df['StartDate'] = pd.to_datetime(df['StartDate'])
df['EndDate'] = pd.to_datetime(df['EndDate'])


(df.assign(dummy=1)
   .merge(cycledates.assign(dummy=1), on='dummy')
   .query('StartDate<=EvaluationDate<EndDate')
   .groupby(['EvaluationDate', 'Customer'])      # ingroupingvars
   ['Apples'].sum()                              #
   .unstack('Customer', fill_value=0)            # remove fill_value=0 if you want the NaN
)

输出:

Customer         A   B  C
EvaluationDate           
2020-01-01       7   6  0
2020-01-02      16   0  4
2020-01-03      12   8  4
2020-01-04       2  17  7
2020-01-05       7  17  3

相关问题 更多 >