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

Question

考虑数据框

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

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