Pandas Conditional Shift Based On Another Column

Question

I have a simple dataframe of people and value. I want to add a column which returns the users next value. This is simple to do with a shift(-1), but I would like the new column to be a null value if the person changes. I could iterate over the rows but that seems like a slow option.

Original dataframe:

import pandas as pd
df = pd.DataFrame(data=[['A', 1], ['A', 4],
                        ['B', 1], ['B', 4]], 
                  columns=['person', 'val'])

    person  val
0   A       1
1   A       4
2   B       1
3   B       4

Desired Output:

    person  val conditional shift
0   A       1   4.0
1   A       4   NaN
2   B       1   6.0
3   B       6   NaN

score 1 · Accepted Answer · answered Apr 04 '20 at 19:37

Looks like you want GroupBy.shift:

df['conditional shift'] = df.groupby('person').val.shift(-1)

print(df)

  person  val  conditional shift
0      A    1                4.0
1      A    4                NaN
2      B    1                6.0
3      B    6                NaN

Pandas Conditional Shift Based On Another Column

1 Answers1