如何按id分组Pandas中含列表的列以生成嵌套列表？

问题场景 #

原始DataFrame：

import pandas as pd
df = pd.DataFrame({'id':[0,0,1], 'lst':[["a", "b", "c"],["e"],["c"]]})
# 输出：
#    id         lst
# 0   0  [a, b, c]
# 1   0        [e]
# 2   1        [c]

期望按id分组后得到嵌套列表的结果：

id        lst
0   0  [[a, b, c], [e]]
1   1        [[c]]

实现方案 #

直接使用groupby配合agg(list)即可实现需求，代码如下：

result = df.groupby('id')['lst'].agg(list).reset_index()
print(result)

原理说明 #

• groupby('id')按id字段分组，将相同id的行归为一组
• agg(list)对每个分组内的lst列执行聚合操作：把分组里的所有列表元素直接收集到一个新列表中，保留原列表的嵌套结构
• reset_index()将分组后的索引重置为常规序列索引，恢复DataFrame的标准格式

关于之前方案失效的说明 #

你尝试的方案大概率是针对列表元素拼接合并的场景（比如把[a,b,c]和[e]合并成[a,b,c,e]），而不是保留原列表的嵌套结构。我们这里直接用list作为聚合函数，正好匹配你需要的嵌套列表输出。

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