Python Pandas: Группировать по совокупной сумме, но избегая суммы, где флаг равен 0

Question

Уже есть несколько вопросов по этому топу c (например, Pandas: совокупная сумма одного столбца на основе значения другого), однако ни один из них не соответствует моим требованиям. Скажем, у меня есть такой фрейм данных:

id  flag
a   1
a   1
a   0
a   0
a   1
b   0
b   0
b   1
b   1
b   1
b   1
c   0
c   1
c   1
c   0
c   1

Я хочу вычислить кумулятивную сумму группировки флагов по идентификатору, но избегаю суммы, когда флаг равен 0, а также снова сбрасывает cumsum до 0. Я попытался подвести итог, используя shift(), groupby(id)['flag']cumsum() внутри np.where, но безуспешно. Желаемый результат должен быть:

id  flag    cum_flag
a   1   1
a   1   2
a   0   0
a   0   0
a   1   1
b   0   0
b   0   0
b   1   1
b   1   2
b   1   3
b   1   4
c   0   0
c   1   1
c   1   2
c   0   0
c   1   1

DDL для генерации фрейма данных:

df = pd.DataFrame({'id': [a, a, a, a, a, b, b, b, b, b, b, c, c, c, c, c],
                 'flag': [1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1]})

Спасибо за вашу помощь!

Answer 1

Вы можете использовать cumsum(), чтобы определить последовательные блоки 1, а затем groupby:

df['cum_flag'] = df.groupby(['id',(df['flag']==0).cumsum()]).cumsum()

Вывод:

   id  flag  cum_flag
0   a     1         1
1   a     1         2
2   a     0         0
3   a     0         0
4   a     1         1
5   b     0         0
6   b     0         0
7   b     1         1
8   b     1         2
9   b     1         3
10  b     1         4
11  c     0         0
12  c     1         1
13  c     1         2
14  c     0         0
15  c     1         1