问题描述
我的数据是这样的：
df = pd.DataFrame({'ym': [202012,202009,202012,202101,202103,202104,202102,202105,202011,202104],'user_id': ['049','009','130','166','282','282']})
df
ym user_id
0 202012 049
1 202009 009
2 202012 009
3 202101 009
4 202103 009
5 202104 009
6 202012 130
7 202101 130
8 202102 130
9 202103 130
10 202105 130
11 202009 166
12 202012 166
13 202101 166
14 202011 282
15 202104 282
我要数数 每个用户的年月假出现，我尝试了这种讨厌的方式来获得：
#get a shift column of current year month column
df['ym_s'] = df.sort_values(['user_id','ym']).groupby('user_id')['ym'].shift(-1)
df = df.sort_values(['user_id','ym'])
#get diff value of year month
df['cntn_vl'] = df['ym_s']- df['ym']
#pick out the consecutive part of in my data: normal consecutive month = 1 or dec-jan value
df['rank'] = np.where(df['cntn_vl'].isin([1,89]),1,0)
#finally aggregate the consecutive appearance of each user
df_1 = df.groupby(by =['user_id'],as_index=False)['rank'].sum()
df_1
user_id rank
0 009 2
1 049 0
2 130 3
3 166 1
4 282 0
我的问题是，有没有更好/更快（在我的数十亿行数据中，很难运行）的方法来解决它。
解决方法
df = df.sort_values(['user_id','ym'])
df['rank'] = df.groupby('user_id')['ym'].diff().isin([1,89])
#finally aggregate the consecutive appearance of each user
df_1 = df.groupby('user_id',as_index=False)['rank'].sum()
或自定义 lambda 函数以避免重复 groupby:
df = df.sort_values(['user_id','ym'])
f = lambda x: x.diff().isin([1,89]).sum()
#numpy alternative should be faster
#f = lambda x: np.in1d(np.diff(x),[1,89]).sum()
df_1 = df.groupby('user_id')['ym'].agg(f).reset_index(name='rank')
print (df_1)
user_id rank
0 009 2
1 049 0
2 130 3
3 166 1
4 282 0