CODEWITHSUNDEEP
pythonpandasshift
elPastor
elPastor

Reputation: 8956

Pandas: Shift one column by other column value

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.

Upvotes: 3

Views: 3020

Answers (2)

MaxU - stand with Ukraine
MaxU - stand with Ukraine

Reputation: 210832

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

Upvotes: 2

elPastor
elPastor

Reputation: 8956

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.

Upvotes: 1

Related Questions