在 Pandas 中按組獲取連續發生的事件

Question

我正在使用一個具有id 、 wage和date的 DataFrame ，如下所示：

id   wage   date
1    100    201212
1    100    201301             
1     0     201302
1     0     201303
1    120    201304
1     0     201305
      .
2     0     201302
2     0     201303

我想創建一個n_months_no_income列來計算給定個人連續獲得多少wage==0 ，如下所示：

id   wage   date     n_months_no_income
1    100    201212             0
1    100    201301             0
1     0     201302             1
1     0     201303             2
1    120    201304             0
1     0     201305             1
      .                        .
2     0     201302             1
2     0     201303             2

我覺得這是groupby('id') ， cumcount() ，也許是diff()或apply()之間的某種混合，然后是fillna(0) ，但我沒有找到合適的。

你有什么想法？

以下是數據框的示例，以便於復制：

df = pd.DataFrame({'id':[1,1,1,1,1,1,2,2],'wage':[100,100,0,0,120,0,0,0],
 'date':[201212,201301,201302,201303,201304,201305,201302,201303]})

編輯：添加了易於使用的代碼。

Answer 1

在您的情況下，使用cumcount兩個groupby並使用cumcount創建cumsum密鑰

df.groupby('id').wage.apply(lambda x : x.groupby(x.ne(0).cumsum()).cumcount())
Out[333]: 
0    0
1    0
2    1
3    2
4    0
5    1
Name: wage, dtype: int64