如何将性别预测函数结果作为新列添加到DataFrame？

问题描述 #

我在含字符串型观测值（客户姓名）的数据集上训练了一个决策树模型（未做文本分词处理），现在要在新数据集customers上测试模型，将预测的性别列追加到这个DataFrame中。目前编写的genderpredictor函数仅能打印结果，无法直接存入DataFrame，需要解决这个问题。

现有预测函数：

def genderpredictor(a):
    test_name1 = [a]
    transform_dv = dv.transform(features(test_name1))
    vector = transform_dv.toarray()
    if dclf.predict(vector) == 0:
        print("Female")
    else:
        print("Male")

现有测试代码（仅输出打印结果）：

customers_list = customers.Cust_First_Name.tolist()
for n in customers_list:
    print(genderpredictor(n))

customers DataFrame结构示例：

解决步骤 #

1. 修改genderpredictor函数，返回结果而非打印 #

原函数用print输出结果，无法被DataFrame捕获，需要改为返回字符串：

def genderpredictor(a):
    test_name1 = [a]
    transform_dv = dv.transform(features(test_name1))
    vector = transform_dv.toarray()
    # 返回预测结果字符串，替代打印操作
    return "Female" if dclf.predict(vector) == 0 else "Male"

2. 用apply方法生成新列并追加到DataFrame #

直接对Cust_First_Name列应用修改后的函数，生成新列Predicted_Gender：

# 生成预测列并追加到原DataFrame
customers['Predicted_Gender'] = customers['Cust_First_Name'].apply(genderpredictor)

# 查看更新后的DataFrame
customers.head()

执行后，customers会新增预测列，示例输出：

补充说明 #

• apply是pandas处理列级批量操作的高效方式，比手动循环更简洁且性能更优；
• 确保dv（特征转换器）和dclf（训练好的决策树模型）已在当前代码环境中正确加载，否则会触发未定义错误。

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