Pandas的groupby.head(-n) 方法会丢弃一些分组。

Question

我有一个如下所示的数据框：

数据框示例

对于每个reviewerID，我希望选取根据reviewTime排序后的最后两个实例作为测试数据，其余的（直到倒数第二个实例）作为训练数据集。这是我目前使用的代码：

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只有不到两行记录的情况下，Pandas将其所有记录都分配给了测试数据框，而非训练数据框。
我想知道如何解决这个问题，即如何保留那些记录少于2条的reviewerID作为训练数据的一部分，而不是验证数据。考虑到我不清楚每个reviewerID有多少条训练数据。

duffymo · Answer

您可以使用自定义函数配合groupby_apply方法来实现这一操作：

# 首先对您的数据框按时间进行排序
df = df.sort_values("reviewTime")

# 从索引中提取reviewerID作为新的数据框，并赋予唯一的位点标签
rid = df.index.get_level_values('reviewerID').to_frame(index=False)

# 定义一个函数，用于获取每组的前n个元素的索引，这里取前n-2个，确保至少有一个元素在验证集中
min_head = lambda x: x[:min(len(x), -2, key=abs)]

# 应用该函数到分组后的reviewerID上，得到训练集的索引位置
train_idx = rid.groupby('reviewerID').apply(min_head).index.levels[1]

# 通过总体索引减去训练集的索引来得到验证集的索引
val_idx = rid.index.difference(train_idx)

# 根据得到的索引创建训练集和验证集的数据框
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

Pankaj Parkar · Answer

由于您有多个条件需要选取数据，可以使用pandas.core.groupby.DataFrameGroupBy.transform方法构建一个选择行的掩码，并执行布尔索引来选取数据：

# 先按'reviewTime'列对数据进行排序
g = df.sort_values('reviewTime')

# 按'reviewerID'列进行分组，且不对组内排序
g = g.groupby('reviewerID', sort=False)

# 要选取的行数
N = 2

# 获取每个分组的大小并枚举行号
# 对于大小超过N的组，选取最后N行
m1 = g.transform('size').gt(N)
m2 = g.cumcount(ascending=False).lt(N)
m = m1 & m2

# 使用布尔索引选取不满足条件的行作为训练集
df_train = df[~m]
# 使用布尔索引选取满足条件的行作为验证集
df_val = df[m]

示例：

# 原始数据df
                            overall
reviewerID asin reviewTime         
0          0    2024-01-01        0
           1    2024-01-03        1
           2    2024-01-05        2
           3    2024-01-02        3
1          4    2024-01-01        4
           5    202-01-02        5
2          6    2024-01-02        6
           7    2024-01-01        7
           8    2024-01-03        8
3          9    2024-01-01        9

# 训练集df_train
                            overall
reviewerID asin reviewTime         
0          0    2024-01-01        0
           3    2024-01-02        3
1          4    2024-01-01        4
           5    2024-01-02        5
2          7    2024-01-01        7
3          9    2024-01-01        9

# 验证集df_val
                            overall
reviewerID asin reviewTime         
0          1    2024-01-03        1
           2    2024-01-05        2
2          6    2024-01-02        6
           8    2024-01-03        8

中间步骤展示：

# 中间结果df
                            overall  size  cumcount     m1     m2      m
reviewerID asin reviewTime                                              
0          0    2024-01-01        0     4         3   True  False  False
           1    2024-01-03        1     4         1   True   True   True
           2    2024-01-05        2     4         0   True   True   True
           3    2024-01-02        3     4         2   True  False  False
1          4    2024-01-01        4     2         1  False   True  False
           5    2024-01-02        5     2         0  False   True  False
2          6    2024-01-02        6     3         1   True   True   True
           7    2024-01-01        7     3         2   True  False  False
           8    2024-01-03        8     3         0   True   True   True
3          9    2024-01-01        9     1         0  False   True  False