Векторизованный расчет значения столбца на основе предыдущего значения этого же столбца?

У меня есть пандас dataframe с 2 столбцами, как это:

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

Я хочу создать новый столбец C следующим образом: C[i]=C[i-1]-A[i]+B[i]

В этом вопросе ответ предлагает использовать цикл вроде этого:

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

Который делает работу.

Но поскольку петли обычно медленнее по сравнению с векторизованными вычислениями, мне было интересно, есть ли в пандах векторизованное решение для этого. (И это причина этого нового вопроса).

Я пытался использовать метод сдвига, как это

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 ответ

Решение

Ваша проблема может быть векторизована, так как delta[i] = C[i] - C[i-1] = -A[i]+B[i]. Сначала вы можете получить дельту из A и B и рассчитать кумулятивную сумму дельты (плюс C[0]), чтобы получить полный C, следующим образом:

In [21]:
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
Другие вопросы по тегам