pandas中如何按userId分组提取对应movieId生成嵌套列表

不用写逐行循环，pandas内置的分组聚合就能高效实现这个需求，性能比Python层写for循环高两个量级以上，代码也更简洁。

核心代码如下：

# 按userId分组，聚合每个用户对应的movieId为列表，最后转为嵌套列表
mv_lst = test_df.groupby('userId')['movieId'].agg(list).tolist()

运行后输出的结果和预期完全一致：

[[169, 2471, 48516], [2571, 109487, 112552, 112556], [356, 2394, 2431]]

补充说明 #

• 如果不需要输出结果按userId升序排列，可以给groupby加参数sort=False，执行速度会更快，返回顺序对应用户在原数据中第一次出现的顺序：

  mv_lst = test_df.groupby('userId', sort=False)['movieId'].agg(list).tolist()
  

• 你之前写的循环逻辑有本质问题：逐行判断当前行的userId是否在全量用户唯一值集合里，这个判断永远成立，因为每一行的userId本来就属于全量用户集合，最终只会把所有movieId平铺成一维列表，完全没有做按用户拆分聚合的动作，再加上逐行遍历pandas对象本身就是性能最差的用法，数据量变大之后运行慢是必然的。

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