快速获取Pandas DataFrame行/列中非空值索引的高效方法

给定如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    'a': {0: 3.0, 1: 2.0, 2: None}, 
    'b': {0: 10.0, 1: None, 2: 8.0}, 
    'c': {0: 4.0, 1: 2.0, 2: 6.0}
})

输出为：

a     b    c
0  3.0  10.0  4.0
1  2.0   NaN  2.0
2  NaN   8.0  6.0

原需求是生成新DataFrame，每行对应原DataFrame该行所有非NaN值的列名。原方法使用apply逐行处理：

df2 = df.apply(lambda x: pd.Series(x.dropna().index), axis=1)

输出符合预期，但apply本质是逐行执行Python循环，在大数据集下性能较差。以下是两种更高效的实现方式：

方法1：利用stack()结合分组（推荐） #

stack()是Pandas的向量化操作，性能远优于逐行apply：

# 堆叠非空值，保留原始行索引和列名
stacked = df.stack().reset_index(level=1)
# 按行索引分组，收集列名并展开为DataFrame
df2 = stacked.groupby(level=0)['level_1'].apply(list).apply(pd.Series)

执行后得到的结果和原方法一致：

0  1    2
0  a  b    c
1  a  c  NaN
2  b  c  NaN

方法2：基于NumPy布尔矩阵的列表推导 #

通过NumPy直接操作布尔矩阵，避免Pandas的逐行循环开销：

# 获取非空值的布尔矩阵
mask = df.notna().to_numpy()
# 逐行提取对应列名，生成列表
col_lists = [df.columns[row_mask].tolist() for row_mask in mask]
# 转换为DataFrame，自动补全NaN到最长行的长度
df2 = pd.DataFrame(col_lists)

关于获取每列非空值的行索引 #

如果需要获取每列的非空值行索引，只需转置原DataFrame后使用上述任意一种方法即可，例如：

# 转置后用stack方法处理
df_T = df.T.stack().reset_index(level=1)
col_non_null_rows = df_T.groupby(level=0)['level_1'].apply(list).apply(pd.Series)

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