Joe Joe - 4 months ago 42
Python Question

Variable shift in Pandas

having two columns A and B in a dataframe:

A B
0 1 6
1 2 7
2 1 8
3 2 9
4 1 10


I would like to create a column C. C must have values of B shifted by value of A:

A B C
0 1 6 NaN
1 2 7 NaN
2 1 8 7
3 2 9 7
4 1 10 9


The command:

df['C'] = df['B'].shift(df['A'])


does not work.
Do you have any other ideas?

Answer

This is tricky due to index alignment, you can define a user func and apply row-wise on your df, here the function will perform a shift on the B column and return the index value (using .name attribute to return the index) of the shifted column:

In [134]:    
def func(x):
    return df['B'].shift(x['A'])[x.name]
df['C'] = df.apply(lambda x: func(x), axis=1)
df

Out[134]:
   A   B    C
0  1   6  NaN
1  2   7  NaN
2  1   8  7.0
3  2   9  7.0
4  1  10  9.0