如何按Pandas DataFrame分组标签获取另一列的唯一值列表

原始数据 #

现有包含Brand和color列的DataFrame：

import pandas as pd

df = pd.DataFrame({'Brand': ['VW', 'VW', 'BMW', 'BMW', 'Mercedes'],
                   'color': ['red', 'red', 'red', 'blue', 'black']})

# 输出原始数据
print(df)

输出结果：

Brand  color
0        VW    red
1        VW    red
2       BMW    red
3       BMW   blue
4  Mercedes  black

需求 #

针对Brand列的每个唯一值，获取对应的color列所有唯一值的列表，期望输出：

Brand       colors
0        VW        [red]
1       BMW  [red, blue]
2  Mercedes      [black]

解决方案 #

方法一：分组后直接聚合唯一值 #

通过groupby分组，结合agg和lambda函数提取每个组的唯一color并转为列表：

# 分组聚合
result = df.groupby('Brand')['color'].agg(lambda x: list(x.unique())).reset_index()
# 修改列名
result.columns = ['Brand', 'colors']

print(result)

方法二：先去重再分组聚合（更高效） #

如果数据量较大，先去除重复的Brand-color组合，再分组合并列表，能提升效率：

# 去重后分组
result = df.drop_duplicates().groupby('Brand')['color'].agg(list).reset_index()
result.columns = ['Brand', 'colors']

print(result)

两种方法都能得到目标输出。

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