Variable shift in Pandas

Question

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 1

I'd use help from numpy to avoid the apply

l = np.arange(len(df)) - df.A.values
df['C'] = np.where(l >=0, df.B.values[l], np.nan)
df

   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

---

simple time test

Answer 2

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