向量化/优化行上滚动大熊猫的计算

时间:2019-02-06 00:26:16

标签: python pandas dataframe lambda

我有一个数据框:

df_dict = {
    'sum': np.nan,
    'src1': [5, 1, 1, 5, 1],
    'src2': [2, 6, 2, 4, 1]
}
df = pd.DataFrame(df_dict)

我目前正在通过以下方式更新“ sum”列:

def transform(x):
    row_num = int(x.name)

    previous_sum = 0
    if row_num > 0:
        previous_sum = df.at[row_num-1,'sum']

    src1 = df.at[row_num,'src1']
    src2 = df.at[row_num,'src2']

    df.at[row_num,'sum'] = previous_sum - src2 + src1

df.apply( lambda x: transform(x), axis=1)

这将导致正确的输出:

   sum  src1  src2
0  3.0     5     2
1 -2.0     1     6
2 -3.0     1     2
3 -2.0     5     4
4 -2.0     1     1

问题在于具有许多行的大型数据帧,并且在许多列中重复此过程非常缓慢。

我想对其进行优化,以便在可能的情况下利用一些内置的pandas / numpy矢量化技术,或任何其他更优化的解决方案。

我的示例提供了更好的解决方案

df['sum'] = df['sum'].shift() - df['src1'] + df['src2']

这会导致错误的输出

   sum  src1  src2
0  NaN     5     2
1  NaN     1     6
2  NaN     1     2
3  NaN     5     4
4  NaN     1     1

2 个答案:

答案 0 :(得分:1)

df['sum'] = df['src1'].cumsum() - df['src2'].cumsum()

答案 1 :(得分:0)

似乎您正在执行累加和。

这可以通过cumsum()函数实现

df['sum'] = df['src'].cumsum()

结果

sum     src
0   1   1
1   3   2
2   6   3
3   10  4
4   15  5