Pandas: expand group variable to encompass the first n obervations of the next group in a dataframe-CodePudding

Here is my df with an example of the two columns that I have to work (groupand value) with and an example of the output that I am trying to achieve(output_wanted):

egdf = pd.DataFrame({'group': ['A']*6 ['B']*6 ['C']*5 ['D']*6,
                      'value': list(range(1, 7))*2 list(range(1, 6)) list(range(1, 7)),
                      'output_wanted': ['A']*8 ['B']*6 ['C']*5 ['D']*4)

The idea is that I would like to expand each group to then encompass the first n (in this example two) entries of the next group in the dataframe. I'm lost for ideas on where to begin here. Has anyone got any idea how one might do this? Note that as in the example, the groups are of unequal sizes... Thanks!

CodePudding user response：

Why not use shift and backfill the first two (n) rows?

egdf['output_wanted'] = egdf.group.shift(2).fillna(method='bfill')

Where 2 can be replaced by n of course. If needed DataFrame could be sorted by group first.