根据同一列的先前值对列值进行矢量化计算?

时间:2015-12-27 18:31:30

标签: python pandas vectorization difference

我有一个pandas数据框,其中包含两列AB,如下所示。

我想要一个矢量化解决方案来创建一个新的列C C[i] = C[i-1] - A[i] + B[i]

df = pd.DataFrame(data={'A': [10, 2, 3, 4, 5, 6], 'B': [0, 1, 2, 3, 4, 5]})

>>> df 
     A  B
 0  10  0
 1   2  1
 2   3  2
 3   4  3
 4   5  4
 5   6  5

以下是使用for-loops的solution

df['C'] = df['A']

for i in range(1, len(df)):
    df['C'][i] = df['C'][i-1] - df['A'][i] + df['B'][i] 

>>> df
    A  B   C
0  10  0  10
1   2  1   9
2   3  2   8
3   4  3   7
4   5  4   6
5   6  5   5

......完成这项工作。

但是,由于与矢量化计算相比,循环速度较慢,我希望在pandas中使用矢量化解决方案

我尝试使用shift()方法,如下所示:

df['C'] = df['C'].shift(1).fillna(df['A']) - df['A'] + df['B']

但它没有帮助,因为移位的C列未经计算更新。它保留了原始值:

>>> df['C'].shift(1).fillna(df['A'])
0    10
1    10
2     2
3     3
4     4
5     5

,这会产生错误的结果。

1 个答案:

答案 0 :(得分:7)

这可以通过以下方式进行矢量化:

  • delta[i] = C[i] - C[i-1] = -A[i] +B[i]。您可以先从deltaA获取B,然后再获取
  • 计算delta(加C[0])的累计总和以获得完整C

代码如下:

delta = df['B'] - df['A']
delta[0] = 0
df['C'] = df.loc[0, 'A'] + delta.cumsum() 
​
print df
    A  B   C
0  10  0  10
1   2  1   9
2   3  2   8
3   4  3   7
4   5  4   6
5   6  5   5