Pandas高效按行提取值不为'Do'的列名（适配大数据量）

给定如下形式的Pandas DataFrame：

A  B  C  D
1 Do  1 All Do
2  5  2 NA Do
3 1  3 NA NA
4 Do  Do  Do  Do

需要按行提取值为"Do"的列名，生成新列Column_names，期望输出如下：

A  B  C  D Column_names
1 Do  1 All Do        A, D
2  5  2 NA Do           D
3 1  3 NA NA        np.nan
4 Do  Do  Do  Do  A, B, C, D

要求实现方法高效，禁止使用apply（避免大数据量下的性能问题），支持动态列数（最多10列）和1000+行数据。

解决方案 #

使用向量化布尔掩码+矩阵点积的方法，完全依赖Pandas/Numpy的底层优化操作，避免逐行循环：

import pandas as pd
import numpy as np

# 1. 构造示例数据（真实场景直接读取数据源即可）
data = {
    'A': ['Do', '5', '1', 'Do'],
    'B': ['1', '2', '3', 'Do'],
    'C': ['All', 'NA', 'NA', 'Do'],
    'D': ['Do', 'Do', 'NA', 'Do']
}
df = pd.DataFrame(data, index=[1,2,3,4])

# 若输入中的"NA"是实际缺失值，先替换为np.nan
# df = df.replace('NA', np.nan)

# 2. 生成布尔掩码：标记每个元素是否等于"Do"
mask = df == "Do"

# 3. 用点积拼接符合条件的列名，去除末尾多余的逗号
df['Column_names'] = mask.dot(df.columns + ',').str.rstrip(',')

# 4. 将空字符串替换为np.nan（对应无"Do"的行）
df['Column_names'] = df['Column_names'].replace('', np.nan)

print(df)

方案说明 #

• 高效性：所有操作都是向量化的，基于Numpy的C底层实现，比apply的Python层面逐行循环快一个数量级以上，完全适配1000+行的大数据量。
• 动态列数支持：代码不依赖固定列名，不管列数是1-10列中的任意数量，都能自动处理。
• 逻辑细节：
  • mask是和原DataFrame形状相同的布尔矩阵，True表示对应位置的值为"Do"。
  • mask.dot(df.columns + ',')通过矩阵点积，将每行中True对应的列名（加逗号）拼接成字符串。
  • str.rstrip(',')去掉每行末尾多余的逗号，避免出现类似"A,B,"的无效格式。
  • 最后将空字符串替换为np.nan，符合示例中无匹配列时的输出要求。

如果你的需求实际是提取值不为"Do"的列名，只需将掩码改为mask = df != "Do"即可，其他代码逻辑保持不变。

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