如何在Pandas中添加按id分组的另一列唯一值列表列

我有如下Pandas DataFrame：

import pandas as pd
df = pd.DataFrame({'id': [1,1,1,1,2,2,2], 'col': ['a','b','c','c','a','b','d']})

输出结果：

id col
0   1   a
1   1   b
2   1   c
3   1   c
4   2   a
5   2   b
6   2   d

希望添加一个新列col2，其中包含按id分组后col列的所有唯一值列表，最终DataFrame如下：

result_df = pd.DataFrame({'id': [1,1,1,1,2,2,2], 'col': ['a','b','c','c','a','b','d'],
             'col2': [['a','b','c'],['a','b','c'],['a','b','c'],['a','b','c'],
                     ['a','b','d'],['a','b','d'],['a','b','d']]})

输出结果：

id col      col2
0   1   a  [a, b, c]
1   1   b  [a, b, c]
2   1   c  [a, b, c]
3   1   c  [a, b, c]
4   2   a  [a, b, d]
5   2   b  [a, b, d]
6   2   d  [a, b, d]

方法一：使用groupby + transform（推荐） #

直接对分组后的col列提取唯一值列表，transform会自动将结果广播到原DataFrame的每一行：

import pandas as pd

df = pd.DataFrame({'id': [1,1,1,1,2,2,2], 'col': ['a','b','c','c','a','b','d']})
df['col2'] = df.groupby('id')['col'].transform(lambda x: list(x.unique()))

方法二：先聚合再合并 #

先分组聚合得到每个id对应的唯一值列表，再通过合并操作关联到原DataFrame：

import pandas as pd

df = pd.DataFrame({'id': [1,1,1,1,2,2,2], 'col': ['a','b','c','c','a','b','d']})
# 生成id与唯一值列表的映射表
id_unique_map = df.groupby('id')['col'].agg(lambda x: list(x.unique())).reset_index(name='col2')
# 合并到原DataFrame
df = df.merge(id_unique_map, on='id', how='left')

两种方法都能实现需求，第一种写法更简洁，无需额外的合并步骤。

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