Python – Pandas: Move one column to another

I’m trying to use the value of one column to move the value of another column by that amount. According to the docs, Pandas shift() accepts an integer, but is there a way to use series instead?

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'])

… This will certainly not work.

Expected 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 were easier, the shift would always be the same, so theoretically the 'b' series could be True/False or some other binary trigger, and .shift() could still be an integer. Going that way feels a bit corny, but it gets the job done.

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