基于行条件为Pandas DataFrame添加分类列

这里有几种简洁的实现方式，完全不用循环：

方法1：字典映射（推荐，大数据集效率更高） #

先构建元素到分类的映射字典，再用map方法批量映射：

import pandas as pd

# 示例数据
df = pd.DataFrame({'Element': ['a_1', 'a_2', 'b_1', 'a_3', 'b_2']})
As = ['a_1', 'a_2', 'a_3']
Bs = ['b_1', 'b_2']

# 构建映射字典（Python3.9+支持|合并字典）
category_map = {elem: 'A' for elem in As} | {elem: 'B' for elem in Bs}
# 兼容旧Python版本的写法：
# category_map = {elem: 'A' for elem in As}
# category_map.update({elem: 'B' for elem in Bs})

df['Category'] = df['Element'].map(category_map)

方法2：isin结合numpy.where #

利用isin判断元素是否属于A列表，再用np.where直接赋值：

import pandas as pd
import numpy as np

df = pd.DataFrame({'Element': ['a_1', 'a_2', 'b_1', 'a_3', 'b_2']})
As = ['a_1', 'a_2', 'a_3']
Bs = ['b_1', 'b_2']

df['Category'] = np.where(df['Element'].isin(As), 'A', 'B')

这两种方法都比循环高效得多，尤其是处理大型数据集时优势明显。其中字典映射的方式扩展性更好——如果后续增加更多分类，只需扩展字典即可。

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