可以在此代码逻辑上使用Apply函数或矢量化吗?

时间:2018-11-27 08:58:54

标签: python pandas numpy vectorization apply

我正在尝试计算期末余额

输入数据框:

    open   inOut    close
0   3      100      0
1   0      300      0
2   0      200      0
3   0      230      0
4   0      150      0

输出数据框

    open    inOut   close
0   3       100     103
1   103     300     403
2   403     200     603
3   603     230     833
4   833     150     983  

我可以使用 for循环来实现此目标,并对其进行优化,我使用了 iterrow()

循环

%%timeit
for i in range(len(df.index)):
    if i>0:
        df.iloc[i]['open'] = df.iloc[i-1]['close']
        df.iloc[i]['close'] = df.iloc[i]['open']+df.iloc[i]['inOut']
    else:
        df.iloc[i]['close'] = df.iloc[i]['open']+df.iloc[i]['inOut'] 

1.64 ms ± 51.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

增长

%%timeit
for index,row in dfOg.iterrows():
    if index>0:
        row['open'] = dfOg.iloc[index-1]['close']
        row['close'] = row['open']+row['inOut']
    else:
        row['close'] = row['open']+row['inOut']

627 µs ± 28.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
  

性能从1.64ms-> 627µs优化

根据此blog,我正在努力找出如何使用apply()和矢量化编写上述逻辑的方法。
对于矢量化,我尝试移动列,但无法实现所需的输出。

2 个答案:

答案 0 :(得分:3)

编辑:我进行了一些更改,以匹配OP对问题所做的编辑

您可以以向量化的方式完成您想做的事情,而无需像这样的循环:

import pandas as pd

d = {'open': [3] + [0]*4, 'inOut': [100, 300, 200, 230, 150], 'close': [0]*5}
df = pd.DataFrame(d)

df['close'].values[:] = df['open'].values[0] + df['inOut'].values.cumsum()
df['open'].values[1:] = df['close'].values[:-1]

使用%%timeit计时:

529 µs ± 5.39 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

输出:

   close  inOut  open
0    103    100     3
1    403    300   103
2    603    200   403
3    833    230   603
4    983    150   833

因此,以这种方式对您的代码进行向量化确实更快一些。实际上,它可能尽可能快。您可以通过仅计时数据帧创建代码来看到此信息:

%%timeit
d = {'open': [3] + [0]*4, 'inOut': [100, 300, 200, 230, 150], 'close': [0]*5}
df = pd.DataFrame(d)

结果:

367 µs ± 5.67 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

减去创建数据帧所需的时间,填充数据帧的矢量化版本仅需约160 µs。

答案 1 :(得分:0)

您可以使用np.where

%%timeit
df['open'] = np.where(df.index==0, df['open'], df['inOut'].shift())
df['close'] = df['open'] + df['inOut']
# 1.07 ms ± 16.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

输出:

    open    inOut   close
0   3.0     100     103.0
1   100.0   300     300.0
2   300.0   200     200.0
3   200.0   230     230.0
4   230.0   150     150.0