矢量化计算Pandas Dataframe

时间:2018-03-07 23:15:05

标签: python pandas numpy

我有一个小问题,我已经使用循环解决了,但我试图看看是否有一种方法可以尝试对其中一些进行矢量化以尝试提高性能。

基本上我有2个数据帧(DF_A和DF_B),其中DF_B中的行基于DF_A中相应行的行和DF_B中的上行行。我确实在DF_B中有第一行值。

df_a = [
  [1,2,3,4]
  [5,6,7,8]
  [..... more rows]
]
df_b = [
 [1,2,3,4] 
 [ rows of all 0 values here, so dimensions match df_a]
]

我想要实现的是df_b中的第二行例如是df_b中第一行的值+ df_a中第二行的值。所以在这种情况下:

df_b.loc[2] = [6,8,10,12] 

我能够使用df_a范围内的循环完成此操作,保持先前的行值保存,然后将当前索引的行添加到前一行值。看起来效率不高。

2 个答案:

答案 0 :(得分:2)

这是一个numpy解决方案。这应该比pandas循环快得多,特别是因为它通过numba使用JIT编译。

from numba import jit

a = df_a.values
b = df_b.values

@jit(nopython=True)
def fill_b(a, b):
    for i in range(1, len(b)):
        b[i] = b[i-1] + a[i]
    return b

df_b = pd.DataFrame(fill_b(a, b))

#     0   1   2   3
# 0   1   2   3   4
# 1   6   8  10  12
# 2  15  18  21  24
# 3  28  32  36  40
# 4  45  50  55  60

效果基准

import pandas as pd, numpy as np
from numba import jit

df_a = pd.DataFrame(np.arange(1,1000001).reshape(1000,1000))

@jit(nopython=True)
def fill_b(a, b):
    for i in range(1, len(b)):
        b[i] = b[i-1] + a[i]
    return b

def jp(df_a):

    a = df_a.values
    b = np.empty(df_a.values.shape)
    b[0] = np.arange(1, 1001)

    return pd.DataFrame(fill_b(a, b))

%timeit df_a.cumsum()  # 16.1 ms
%timeit jp(df_a)       # 6.05 ms

答案 1 :(得分:1)

您可以使用累计金额df_b创建df_a,就像这样

df_a = pd.DataFrame(np.arange(1,17).reshape(4,4))
df_b = df_a.cumsum()

    0   1   2   3
0   1   2   3   4
1   6   8  10  12
2  15  18  21  24
3  28  32  36  40