Pandas groupby.head(-n) 删除一些组

Question

我有一个如下所示的 df：

对于每个 reviewerID，我想选择最后 2 个实例（按 reviewTime 排序）作为测试数据，其余的（最多最后 2 个实例）作为训练数据集。这是我的代码：

df = df.sort_values("reviewTime")
df_train = df.groupby('reviewerID').head(-2).reset_index(drop=True)
df_val = df.groupby('reviewerID').tail(2).reset_index(drop=True)

但是，此代码从训练数据中删除了一些组（reviewerID）。因此，如果我获得原始数据帧和 df_train 中唯一 reviewerID 的数量，我将得到不同的数字：

print(df_train['reviewerID'].nunique())
print(df_val['reviewerID'].nunique())
print(df['reviewerID'].nunique())

776775
777242
777242

请注意，测试集具有正确的值。 现在我猜测，如果我的 reviewerID 少于 2 行，pandas 会将所有这些分配给测试 df 而不是训练 df。 我现在想知道如何解决这个问题，并将小于 2 个值的 reviewerID 作为训练数据而不是验证的一部分。考虑到我不知道每个 reviewerID 有多少训练数据。

Answer 1

您可以使用自定义函数

groupby_apply

:

# Sort your dataframe
df = df.sort_values("reviewTime")

# Create a dataframe from reviewerID with unique position labels
rid = df.index.get_level_values('reviewerID').to_frame(index=False)

# Extract index positions for df_train & df_val
min_head = lambda x: x[:min(len(x), -2, key=abs)]
train_idx = rid.groupby('reviewerID').apply(min_head).index.levels[1]
val_idx = rid.index.difference(train_idx)

# Create datasets
df_train = df.iloc[train_idx]
df_val = df.iloc[val_idx]

输出：

>>> df_train.sort_index()
                              overall
reviewerID asin   reviewTime         
0.0        661.0  2015-08-12        1
           959.0  2015-08-13        2
           2350.0 2015-08-14        3
1.0        458.0  2015-08-12        1

>>> df_val.sort_index()
                              overall
reviewerID asin   reviewTime         
0.0        2807.0 2015-08-15        4
           9087.0 2015-08-16        5

>>> df.sort_index()
                              overall
reviewerID asin   reviewTime         
0.0        661.0  2015-08-12        1
           959.0  2015-08-13        2
           2350.0 2015-08-14        3
           2807.0 2015-08-15        4
           9087.0 2015-08-16        5
1.0        458.0  2015-08-12        1