I would like to vectorize this piece of python code with for loop conditioned on current state for speed and efficiency.
values for df_B are computed based on current-state (state
) AND corresponding df_A value.
Any ideas would be appreciated.
JavaScript
x
17
17
1
import pandas as pd
2
df_A = pd.DataFrame({'a': [0, 1, -1, -1, 1, -1, 0, 0] ,})
3
df_B = pd.DataFrame( data=0, index=df_A.index, columns=['b'])
4
print(df_A)
5
6
state = 0
7
for index, iter in df_A.iterrows():
8
if df_A.loc[index ,'a'] == -1:
9
df_B.loc[index ,'b'] = -10 -state
10
elif df_A.loc[index, 'a'] == 1:
11
df_B.loc[index, 'b'] = 10 - state
12
elif df_A.loc[index, 'a'] == 0:
13
df_B.loc[index, 'b'] = 0 - state
14
temp_state = state
15
state += df_B.loc[index, 'b']
16
print(df_B)
17
Advertisement
Answer
This seems overkill. Your state
variable basically is the previous value in df_A['a']*10
. So we can just use shift
:
JavaScript
1
4
1
s = df_A['a'].mul(10)
2
3
df_B['b'] = s - s.shift(fill_value=0)
4