Pandas groupby.head(-n)丢失分组：将不足2条数据的reviewerID归入训练集

问题分析 #

你当前的代码里，groupby.head(-2)会自动过滤掉那些组内数据不足2条的用户——因为head(-2)是取每组前n-2条数据，当组内数据量n<2时，这部分就没有数据返回，所以这些用户的所有数据都被归入了tail(2)的验证集里，导致训练集的唯一用户数比原数据集少，验证集却包含了所有用户。

现在需要调整逻辑：只有当用户的记录数≥2条时，才把最后2条放入验证集；记录数不足2条的用户，所有数据全部归入训练集。

解决方案代码 #

# 先按用户分组，组内按reviewTime排序
df = df.sort_values(["reviewerID", "reviewTime"])
# 给每个用户的每条记录按顺序编序号（从1开始计数）
df["row_num"] = df.groupby("reviewerID").cumcount() + 1
# 给每条记录标记所属用户的总记录数
df["total_rows"] = df.groupby("reviewerID")["row_num"].transform("max")

# 拆分训练集：要么用户记录不足2条（全留），要么取除最后2条外的所有记录
df_train = df[(df["total_rows"] < 2) | (df["row_num"] <= df["total_rows"] - 2)].reset_index(drop=True)
# 拆分验证集：仅保留记录数≥2的用户的最后2条记录
df_val = df[(df["total_rows"] >= 2) & (df["row_num"] > df["total_rows"] - 2)].reset_index(drop=True)

# 可选：清理临时生成的辅助列
df = df.drop(["row_num", "total_rows"], axis=1)

代码说明 #

1. 排序与编号：先按reviewerID分组，组内按reviewTime排序，再给每个用户的记录编序号，这样能精准定位每组的最后2条数据。
2. 标记总记录数：用transform方法让每条记录都能获取到所属用户的总记录数，方便后续筛选。
3. 筛选逻辑：
   • 训练集涵盖两类数据：记录数不足2条的用户的全部数据，以及记录数≥2的用户除最后2条外的所有数据。
   • 验证集只保留记录数≥2的用户的最后2条数据，彻底排除了记录数不足2条的用户。

结果验证 #

执行后可以通过以下方式确认效果：

• 训练集的唯一reviewerID数量会和原数据集完全一致（777242）
• 验证集的唯一reviewerID数量等于原数据集中记录数≥2的用户数量
• 训练集与验证集的总记录数之和等于原数据集的总记录数

内容的提问来源于stack exchange，提问作者Zeinab Sobhani