Pandas: number rows within group cumulatively and across another group-CodePudding

Given the following dataframe:

    col_1 col_2 col_3
0     1     A     1
1     1     B     1
2     2     A     3
3     2     A     3
4     2     A     3
5     2     B     3
6     2     B     3
7     2     B     3
8     3     A     2
9     3     A     2
10    3     C     2
11    3     C     2

I need to create a new column in which the rows are numbered cumulatively within each group formed by 'col_1' and 'col_2', but also cumulatively after each group of 'col_1', like this:

    col_1 col_2 col_3  new
0     1     A     1     1
1     1     B     1     1
2     2     A     3     2
3     2     A     3     3
4     2     A     3     4
5     2     B     3     2
6     2     B     3     3
7     2     B     3     4
8     3     A     2     5
9     3     A     2     6
10    3     C     2     5
11    3     C     2     6

I've tried:

df['new'] = df.groupby(['col_1', 'col_2']).cumcount()   1

But this doesn't add up from the previous group as intended.

Thank you in advance!

CodePudding user response：

You don't want to reset on 'col_1', so don't group on it!

df['new'] = df.groupby('col_2').cumcount()   1

output:

    col_1 col_2  col_3  new
0       1     A      1    1
1       1     B      1    1
2       2     A      3    2
3       2     A      3    3
4       2     A      3    4
5       2     B      3    2
6       2     B      3    3
7       2     B      3    4
8       3     A      2    5
9       3     A      2    6
10      3     C      2    1
11      3     C      2    2

CodePudding user response：

This is a tricky problem. You want to calculate the cumcount within group, but for all subsequent groups you need to keep track of how much that was already incremented so you know the offset to apply. That can be done with a max cumsum of this cumcount over the previous groups. Here the only complication is that you need to determine the relationship between previous and subsequent group labels, in case there isn't some simple 1 increment between groups.

# Cumcount within group
s = df.groupby(['col_1', 'col_2']).cumcount()

# Determine how many cumcounts were within all previous groups of `col_1' 
to_merge = s.add(1).groupby(df['col_1']).max().cumsum().add(1).to_frame('new')

# Link group with prior group label
df1 = df[['col_1']].drop_duplicates()
df1['col_1_shift'] = df1['col_1'].shift(-1)
df1 = pd.concat([to_merge, df1.set_index('col_1')], axis=1)

# Bring the group offset over
df = df.merge(df1, left_on='col_1', right_on='col_1_shift', how='left')

# Add the group offset to the cumulative count within group.
# First group (no previous group) is NaN so fill with 1.
df['new'] = df['new'].fillna(1, downcast='infer')   s

# Clean up merging column
df = df.drop(columns='col_1_shift')

    col_1 col_2  col_3  new
0       1     A      1    1
1       1     B      1    1
2       2     A      3    2
3       2     A      3    3
4       2     A      3    4
5       2     B      3    2
6       2     B      3    3
7       2     B      3    4
8       3     A      2    5
9       3     A      2    6
10      3     C      2    5
11      3     C      2    6