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