Search code examples
pythonpandasshift

how to shift non nan value in multiple columns row wise by group? (2nd)

i have post a similar question which was solved by jezrael perfertly. but this time i have the same dataframe with a new row called label as bellow:

A1 A2 A3 A4 A5 A6 label
1 nan 3 7 nan 8 A
nan 5 nan 11 9 nan A
54 6 84 12 3 nan A
10 nan nan 16 nan 45 B
12 93 13 31 5 91 B
73 nan 45 nan nan 9 B

i want to shift non-nan value n rows according to the label column.

desire output for n = 1

A1 A2 A3 A4 A5 A6 label
nan nan nan nan nan nan A
nan nan nan 7 nan nan A
1 5 3 11 9 nan A
nan nan nan nan nan nan B
10 nan nan 16 nan 45 B
12 nan 13 nan nan 91 B

the solution in previous post without label column is

df = df.apply(lambda x: x.dropna().shift(1))

so i tried

columns = df.drop(columns = ['label']).columns
df[columns] = df.groupby(['label'])[columns].apply(lambda x: x.dropna().shift(1))

which it only leave rows with no nan value in all columns and i can only try using the loop solution

for column in columns:
        df[column] = df.groupby(['label'])[column].apply(lambda x: x.dropna().shift(1)).droplevel(level=0, axis=0)

and again which is slow when number of columns grow large. Wonder if there is anyway to make my attempted solution works.

Solution

  • Let's use transform instead of apply

    df.groupby('label')[columns].transform(lambda s: s.dropna().shift(1))

    Result

         A1   A2    A3    A4   A5    A6
0   NaN  NaN   NaN   NaN  NaN   NaN
1   NaN  NaN   NaN   7.0  NaN   NaN
2   1.0  5.0   3.0  11.0  9.0   NaN
3   NaN  NaN   NaN   NaN  NaN   NaN
4  10.0  NaN   NaN  16.0  NaN  45.0
5  12.0  NaN  13.0   NaN  NaN  91.0

    Some notes:

    • Transform will operate on each column separately as opposed to all columns when using groupby apply
    • Transform will also broadcast the values back to the original shape after dropping the NaN's