Pandas: количество строк внутри группы кумулятивно и по другой группе

Учитывая следующий кадр данных:

 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

Мне нужно создать новый столбец, в котором строки нумеруются кумулятивно в каждой группе, образованной «col_1» и «col_2», а также кумулятивно после каждой группы «col_1», например:

 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

Я пробовал:

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

Но это не складывается из предыдущей группы, как предполагалось.

Решение проблемы

Это сложная проблема. Вы хотите рассчитать cumcount внутри группы, но для всех последующих групп вам нужно отслеживать, сколько уже было увеличено, чтобы вы знали применяемое смещение. Это можно сделать с помощью max+ cumsumнад cumcountпредыдущими группами. Здесь единственная сложность заключается в том, что вам нужно определить отношения между предыдущими и последующими метками групп, в случае, если между метками последующих групп нет простого приращения +1.

# 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