python之在具有条件增量的 pandas 数据框上使用 cumcount

考虑数据框

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

我想找到每个 key 的累积(运行)计数(从 1 开始) ，如果组中的前一个值有 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']) 

请您参考如下方法:

使用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