我有一个pandas数据框,其中包含两列A
,B
,如下所示。
我想要一个矢量化解决方案来创建一个新的列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
,这会产生错误的结果。
答案 0 :(得分:7)
这可以通过以下方式进行矢量化:
delta[i] = C[i] - C[i-1] = -A[i] +B[i]
。您可以先从delta
和A
获取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