python - 删除少于 X 个连续日期的观测值

Question

以下数据框包含同一公司(列 ID)在不同日期(列日期)的数据。我想删除少于 3 天的观察结果。

起始数据集是

df = pd.DataFrame({"ID":{"0":1,"1":1,"2":1,"3":1,"4":4,"5":4,"6":4,"7":2,"8":2,"9":3,"10":3},
    "date":{"0":1421020800000,"1":1421193600000,"2":1422489600000,"3":1423353600000,"4":1421020800000,"5":1421107200000,"6":1421193600000,"7":1421020800000,"8":1421107200000,"9":1421452800000,"10":1421539200000},
    "variable":{"0":28,"1":62,"2":60,"3":72,"4":28,"5":61,"6":62,"7":23,"8":70,"9":32,"10":55}})
df.date = pd.to_datetime(df.date, unit='ms')
df.sort_values(by=["ID", "date"],inplace=True)

在上面的数据框中，只有 ID = 4 的公司满足要求，我想删除其他公司。

我编写了以下代码，但它有一个明显的问题，我不知道如何修复它:

df['delete'] = 0
for name, group in df.groupby(by = "ID"):
    if group.shape[0] < 3:
        df.loc[df['ID']==name,'delete'] = 1
df = df.loc[df['delete'] == 0,:]

上面的代码保留了ID=1和ID=4的两家公司； ID=1 应该被取消，因为它包含 4 个数据点，但其中最多有两个数据点是连续的(而我想强制至少为 3 个)。

任何帮助将不胜感激。谢谢

Answer 1

IIUC 使用 diff + cumsum 和 date 列创建组 key New，然后我们只使用 groupby + 过滤不需要的组

df['New']=df.groupby('ID').date.apply(lambda x : x.diff().dt.days.ne(1).cumsum())
yourdf=df.groupby(['ID','New']).filter(lambda x : len(x)>=3)
yourdf
Out[809]: 
   ID       date  variable  New
4   4 2015-01-12        28    1
5   4 2015-01-13        61    1
6   4 2015-01-14        62    1