在Python中使用isin实现DataFrame三类标签列的生成

问题背景 #

现有如下DataFrame：

# Import pandas library
import pandas as pd
import numpy as np

# initialize list elements
data = ['george',
        'instagram',
        'nick',
        'basketball',
        'tennis']
  
# Create the pandas DataFrame with column name is provided explicitly
df = pd.DataFrame(data, columns=['Unique Words'])

需要基于以下两个列表生成Label列：

key_words = ["football", "basketball", "tennis"]
usernames = ["instagram", "facebook", "snapchat"]

标记规则：

• 属于key_words的词标记为1
• 属于usernames的词标记为2
• 其余词标记为0

最终预期效果：

解决方案 #

方法一：使用numpy.select()（高效推荐）

通过多条件映射实现，适合处理大数据集：

# 定义匹配条件与对应标记值
conditions = [
    df['Unique Words'].isin(key_words),
    df['Unique Words'].isin(usernames)
]
values = [1, 2]

# 生成Label列，未匹配到的内容默认标记为0
df['Label'] = np.select(conditions, values, default=0)

# 输出结果
print(df)

方法二：使用apply()逐行处理

逻辑直观，适合小数据集或需要自定义复杂判断的场景：

def assign_label(word):
    if word in key_words:
        return 1
    elif word in usernames:
        return 2
    else:
        return 0

df['Label'] = df['Unique Words'].apply(assign_label)

两种方法均可得到符合要求的Label列，其中np.select()在数据量较大时性能更优。

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