如何获取DataFrame每行中非NaN值对应的列名（Python）

需求说明 #

给定包含NaN值（列可包含字符串类型）的DataFrame，需要为每行生成对应非NaN值的列名列表/数组。

示例DataFrame #

feature1    feature2    feature3   feature4
  10           NaN          5          2
  2            1            3          1
  NaN          2            4          NaN

预期结果 #

res = [
    ['feature1', 'feature3', 'feature4'],
    ['feature1', 'feature2', 'feature3', 'feature4'],
    ['feature2', 'feature3']
]

解决方案 #

方法一：逐行处理（简洁直观） #

利用apply遍历每行，通过dropna()过滤NaN值后提取列名：

import pandas as pd
import numpy as np

# 构造示例DataFrame
df = pd.DataFrame({
    'feature1': [10, 2, np.nan],
    'feature2': [np.nan, 1, 2],
    'feature3': [5, 3, 4],
    'feature4': [2, 1, np.nan]
})

# 生成目标列表
res = df.apply(lambda row: row.dropna().index.tolist(), axis=1).tolist()
print(res)

方法二：向量化操作（高效适合大数据集） #

先生成非NaN值的布尔掩码，再通过掩码提取对应列名，避免逐行循环：

import pandas as pd
import numpy as np

# 构造示例DataFrame
df = pd.DataFrame({
    'feature1': [10, 2, np.nan],
    'feature2': [np.nan, 1, 2],
    'feature3': [5, 3, 4],
    'feature4': [2, 1, np.nan]
})

# 生成非NaN值的布尔掩码
mask = df.notna()
# 按行提取对应列名
res = [df.columns[mask.iloc[i]].tolist() for i in range(len(df))]
print(res)

方法对比 #

• 方法一代码简洁，逻辑清晰，适合小规模数据集；
• 方法二采用向量化思路，执行效率更高，适合处理大数据量的DataFrame。

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