首页 > 解决方案 > 向量化/优化行上的滚动熊猫计算

问题描述

我有一个数据框:

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

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

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

标签: pythonpandasdataframelambda

解决方案


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


推荐阅读