使用 cumcount 對 Pandas 數據框進行條件增量

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

請注意，基本上每個key組中最后一行的cond值沒有影響。

只是做一個簡單的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

Answer 2

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 3

將groupby與自定義 lambda 函數與先前值的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