如何在Pandas中为列表类型列批量分配类别名称？

我有一个Pandas DataFrame，其中categories列为列表类型，示例数据如下：

categories
[0, 0, 2, 2, 2]
[0, 0, 2, 2]
[0, 0, 2, 2, 2]
[1, 1, 2, 2]
[2, 2, 0, 0]
[1, 0, 2, 3]

样本列表：

li = [[0, 0, 2, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2, 2], [1, 1, 2, 2], [2, 2, 0, 0], [1, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2, 2], [0, 0, 2, 2], [2, 2, 0, 0], [2, 2, 0, 0], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [2, 2, 0, 0], [2, 2], [1, 1, 2], [0, 2, 2, 0], [0, 0, 2, 2], [0, 1], [0, 0], [0, 0, 2, 2], [0, 0], [0, 0, 2, 2], [0, 2, 2, 0], [2, 2, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [2, 2, 0, 0], [0, 0, 2, 2], [2, 2, 0, 1], [2, 2, 0, 0], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2, 2], [2, 1], [0, 0, 2, 2, 2], [2, 2, 0, 0], [2, 0], [2, 2, 0, 0], [0, 2], [0, 2, 2], [0, 0, 2, 2], [0, 2, 2, 0], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2, 2], [0, 0, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [3, 2, 0, 0], [0, 0], [0, 0, 2, 2], [0, 0, 2, 2, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [1, 3], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 0, 2, 2], [0, 2, 0, 2], [0, 0, 2, 2], [2, 2, 0, 0], [2, 2, 0, 0], [2, 2], [0, 0, 2, 2], [0, 2], [0, 0, 2, 2], [0, 0, 2, 2], [2, 2, 0], [2, 2, 0, 0], [0, 0, 2, 2], [0, 0, 2], [2], [0, 0, 2, 2], [2, 2, 2, 1, 1], [0, 0], [0, 3], [2, 2], [1, 2], [1,3]]

需按以下优先级规则生成新列class_name：

• 若列表中同时包含1和3，设为class1；
• 剩余行中，若包含1，设为class2；
• 剩余行中，若包含3，设为class3；
• 剩余行中，若同时包含0和2，设为class4；
• 剩余行中，若包含0，设为class5；
• 剩余行中，若包含2，设为class6。

目前用循环实现效率很低，求无需循环的高效实现方式。

可以利用Pandas的矢量化布尔索引结合集合判断实现，完全避免循环，大幅提升效率。以下是两种实现方式：

方式一：基于集合的快速判断 #

先将每个列表转换为集合（集合的成员判断是O(1)复杂度），再按优先级依次赋值：

import pandas as pd

# 构造DataFrame
df = pd.DataFrame({'categories': li})

# 将列表转换为集合，方便快速判断元素存在性
df['cat_set'] = df['categories'].apply(set)

# 初始化class_name列
df['class_name'] = None

# 按优先级规则赋值
# 规则1：同时包含1和3
mask = df['cat_set'].apply(lambda x: 1 in x and 3 in x)
df.loc[mask, 'class_name'] = 'class1'

# 规则2：剩余行包含1
mask = (df['class_name'].isna()) & df['cat_set'].apply(lambda x: 1 in x)
df.loc[mask, 'class_name'] = 'class2'

# 规则3：剩余行包含3
mask = (df['class_name'].isna()) & df['cat_set'].apply(lambda x: 3 in x)
df.loc[mask, 'class_name'] = 'class3'

# 规则4：剩余行同时包含0和2
mask = (df['class_name'].isna()) & df['cat_set'].apply(lambda x: 0 in x and 2 in x)
df.loc[mask, 'class_name'] = 'class4'

# 规则5：剩余行包含0
mask = (df['class_name'].isna()) & df['cat_set'].apply(lambda x: 0 in x)
df.loc[mask, 'class_name'] = 'class5'

# 规则6：剩余行包含2
mask = (df['class_name'].isna()) & df['cat_set'].apply(lambda x: 2 in x)
df.loc[mask, 'class_name'] = 'class6'

# 可选：删除临时的cat_set列
df = df.drop('cat_set', axis=1)

方式二：预提取元素存在性（更高效） #

如果数据集规模极大，可以先一次性提取各目标元素的存在情况，再组合布尔条件，减少apply调用次数：

import pandas as pd

# 构造DataFrame
df = pd.DataFrame({'categories': li})

# 一次性生成各元素是否存在的布尔列
df['has_1'] = df['categories'].apply(lambda x: 1 in x)
df['has_3'] = df['categories'].apply(lambda x: 3 in x)
df['has_0'] = df['categories'].apply(lambda x: 0 in x)
df['has_2'] = df['categories'].apply(lambda x: 2 in x)

# 初始化class_name列
df['class_name'] = None

# 按优先级赋值
mask = df['has_1'] & df['has_3']
df.loc[mask, 'class_name'] = 'class1'

mask = (df['class_name'].isna()) & df['has_1']
df.loc[mask, 'class_name'] = 'class2'

mask = (df['class_name'].isna()) & df['has_3']
df.loc[mask, 'class_name'] = 'class3'

mask = (df['class_name'].isna()) & df['has_0'] & df['has_2']
df.loc[mask, 'class_name'] = 'class4'

mask = (df['class_name'].isna()) & df['has_0']
df.loc[mask, 'class_name'] = 'class5'

mask = (df['class_name'].isna()) & df['has_2']
df.loc[mask, 'class_name'] = 'class6'

# 删除临时布尔列
df = df.drop(['has_1', 'has_3', 'has_0', 'has_2'], axis=1)

这两种方式都利用了Pandas的矢量化操作，比循环快几个数量级，适合处理大规模数据。

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