在 pandas 数据帧上使用 cumcount 并有条件增量

考虑数据框

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 and 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'])

Use groupby与组合shift and 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