使用 cumcount 对 Pandas 数据框进行条件增量

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前一个元素的值。我如何才能考虑到这一点？

编辑

由于下面的一些解决方案在某些边缘情况下不起作用，我将提供一个更全面的数据框进行测试。

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

使用groupby与组合shift和cumsum。

df['new'] = 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