Pandas: Shift one column by other column value

Question

I'm trying to use one column's values to shift another columns values by that amount. Pandas shift() , per the documentation, takes an integer, but is there a way to instead use a Series?

Current Code:

import pandas as pd

df = pd.DataFrame({ 'a':[1,2,3,4,5,6,7,8,9,10],
                    'b':[0,0,0,0,4,4,4,0,0,0]})

df['a'] = df['a'].shift(df['b'])

...which is of course not working.

Desired output:

    a  b
0   1  0
1   2  0
2   3  0
3   4  0
4   1  4
5   2  4
6   3  4
7   8  0
8   9  0
9  10  0

If it makes it easier, the shift will always be the same, so theoretically the 'b' series could be True / False or some other binary trigger, and the .shift() could still be an integer. Feels a little hacky going that route, but it would get the job done.

Answer 1

we can use numba solution :

from numba import jit

@jit
def dyn_shift(s, step):
    assert len(s) == len(step), "[s] and [step] should have the same length"
    assert isinstance(s, np.ndarray), "[s] should have [numpy.ndarray] dtype"
    assert isinstance(step, np.ndarray), "[step] should have [numpy.ndarray] dtype"
    N = len(s)
    res = np.empty(N, dtype=s.dtype)
    for i in range(N):
        res[i] = s[i-step[i]]
    return res

result:

In [302]: df['new'] = dyn_shift(df['a'].values, df['b'].values)
# NOTE: we should pass Numpy arrays:   ^^^^^^^         ^^^^^^^

In [303]: df
Out[303]:
    a  b  new
0   1  0    1
1   2  0    2
2   3  0    3
3   4  0    4
4   5  4    1
5   6  4    2
6   7  4    3
7   8  0    8
8   9  0    9
9  10  0   10

Answer 2

Figured it out:

df.loc[df['b'] == 4, 'a'] = df['a'].shift(4)

...this is the 'hacky' version I referred to above. The first 4 is really just a trigger and the second 4 would be hard-coded.