Pandas 每组只选择前 3 个 YYYYMM

Question

C下午好，

我有一个像下面这样的数据

+---+---+--------+
|   |USR| MMMMYY |
+---+---+--------+
| 1 | A | 200002 |
+---+---+--------+
| 2 | A | 200003 |
+---+---+--------+
| 3 | A | 200004 |
+---+---+--------+
| 4 | A | 200005 |
+---+---+--------+
| 5 | B | 200001 |
+---+---+--------+
| 6 | B | 200003 |
+---+---+--------+
| 7 | B | 200008 |
+---+---+--------+
| 8 | B | 200009 |
+---+---+--------+

我只需要获得每个 USR 的前三个 *CONSECUTIVE MMMMYY。

+---+---+--------+
|   |USR| MMMMYY |
+---+---+--------+
| 1 | A | 200002 |
+---+---+--------+
| 2 | A | 200003 |
+---+---+--------+
| 3 | A | 200004 |
+---+---+--------+
| 5 | B | 200001 |
+---+---+--------+
| 6 | B | 200003 |
+---+---+--------+

我能够使用 head(3) 获取前 3 条记录

df.sort_values(['USR', 'MMMMYY' ], ascending=[True, True]).groupby('USR', as_index=False).head(3)

但它当然不会带回我需要的东西，也没有使用

df['mm_dif']=df.groupby(['USR'])['MMMMYY'].diff()

df['mm_dif2']=df.groupby(['USR'])['MMMMYY'].diff(-1)

df['check']=np.where((df.mm_dif==1) | (df.mm_dif2==-1),True,False)

当 ['check'] 为真时，它会得到连续的，但在某些情况下，我可能只需要得到 200001 和 200003 并且它们之间不是连续的。任何指导将不胜感激

谢谢

Answer 1

你的MMMMYY是日期时间，然后把它转成datetime先输入：

df['MMMMYY'] = pd.to_datetime(df.MMMMYY, format='%Y%m')

s = df.groupby('USR')['MMMMYY'].transform('min') + pd.offsets.MonthOffset(3)

df[df.MMMMYY<s]

输出：

  USR     MMMMYY
1   A 2000-02-01
2   A 2000-03-01
3   A 2000-04-01
5   B 2000-01-01
6   B 2000-03-01