如何为DataFrame添加包含非零特征名称的新列？

需求说明 #

现有一个包含5000+行的DataFrame，结构示例如下，需要新增一列Values，存储每行中值为1的特征名称（即Power、Security、Humality、Tolerance这些列里值非零的列名），最终保留Id、Context、Values三列。

原始DataFrame示例 #

Id     Context    Power    Security     Humality   Tolerance  

 0       have.       0         0           1           0
 1       has         1         1           1           1
 2       has         0         0           1           1

处理后目标效果 #

Id     Context         Values 

 0       have.        ['Humality']
 1       has         ['Power', 'Security', 'Humality', 'Tolerance']
 2       has         ['Humality', 'Tolerance']

解决方案 #

方法1：简洁直观的逐行处理 #

适合快速实现，代码可读性强：

import pandas as pd

# 假设你的DataFrame名为df
feature_cols = ['Power', 'Security', 'Humality', 'Tolerance']

# 提取每行值为1的特征列名，生成Values列
df['Values'] = df[feature_cols].apply(lambda row: row.index[row == 1].tolist(), axis=1)

# 筛选保留目标列
final_df = df[['Id', 'Context', 'Values']]

方法2：高效矢量化操作（适配5000+行数据集） #

针对大数据量优化，避免逐行循环，处理速度更快：

import pandas as pd

feature_cols = ['Power', 'Security', 'Humality', 'Tolerance']

# 生成布尔矩阵，标记每行各特征是否为1
bool_matrix = df[feature_cols].eq(1)

# 通过矩阵点积拼接符合条件的列名，再分割为列表
df['Values'] = bool_matrix.dot([col + ',' for col in feature_cols]).str.rstrip(',').str.split(',')

# 保留目标列
final_df = df[['Id', 'Context', 'Values']]

执行上述代码后，final_df即为符合要求的结果。

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