Использовать cumcount на pandas dataframe с условным приращением

Question

Рассмотрим фрейм данных

df = pd.DataFrame(
    [
        ['A', 1],
        ['A', 1],
        ['B', 1],
        ['B', 0],
        ['A', 0],
        ['A', 1],
        ['B', 1]
    ], columns = ['key', 'cond'])

Я хочу найти совокупное (текущее) число (начиная с 1) для каждого key, где мы увеличиваем его, только если предыдущее значение в группе имело cond == 1. При добавлении к вышеуказанному фрейму данных это даст

df_result = pd.DataFrame(
    [
        ['A', 1, 1],
        ['A', 1, 2],
        ['B', 1, 1],
        ['B', 0, 2],
        ['A', 0, 3],
        ['A', 1, 3],
        ['B', 1, 2]
    ], columns = ['key', 'cond'])

Обратите внимание, что по существу значения cond последних строк в каждой группе key не действуют.

Просто делаю простые group и cumcount

df.groupby('key').cumcount()

, конечно, не учитывает значение cond предыдущего элемента. Как я могу принять это во внимание?

EDIT

Поскольку некоторые из приведенных ниже решений не работают в некоторых крайних случаях, я приведу более полный набор данных для тестирования.

df = pd.DataFrame(
    [
        ['A', 0],
        ['A', 1],
        ['A', 1],
        ['B', 1],
        ['B', 0],
        ['A', 0],
        ['A', 1],
        ['B', 1],
        ['B', 0]
    ], columns = ['key', 'cond'])

, который при добавлении истинного результата должен дать

df_result = pd.DataFrame(
    [
        ['A', 0, 1],
        ['A', 1, 1],
        ['A', 1, 2],
        ['B', 1, 1],
        ['B', 0, 2],
        ['A', 0, 3],
        ['A', 1, 3],
        ['B', 1, 2],
        ['B', 0, 3]
    ], columns = ['key', 'cond'])

Answer 1

df['cnt'] = df[df["cond"]==1].groupby('key').cumcount()+1
df['cnt'] = df.groupby('key')['cnt'].fillna(method='bfill')
df
# =>   key  cond  cnt
#    0   A     1  1.0
#    1   A     1  2.0
#    2   B     1  1.0
#    3   B     0  2.0
#    4   A     0  3.0
#    5   A     1  3.0
#    6   B     1  2.0

Answer 2

Используйте groupby с пользовательской лямбда-функцией с shift для предыдущих значений, замените первое NaN на обратное заполнение, если возможно, первое значение составляет 0 на клавишу ипоследний cumsum с преобразованием в int:

df['new'] = df.groupby('key')['cond'].apply(lambda x: x.shift().bfill().cumsum()).astype(int)

print (df)
  key  cond  new
0   A     1    1
1   A     1    2
2   B     1    1
3   B     0    2
4   A     0    3
5   A     1    3
6   B     1    2

Answer 3

Используйте groupby с комбинацией shift и cumsum.

df['result'] = df.groupby('key').cond.apply(
        lambda x: x.shift().fillna(1).cumsum()
).astype(int)

df
  key  cond  new
0   A     1    1
1   A     1    2
2   B     1    1
3   B     0    2
4   A     0    3
5   A     1    3
6   B     1    2