如何在pandas DataFrame中获取每行大于0值对应的列索引

核心思路 #

全程使用 numpy + pandas 向量化操作，避免 Python 层的逐行循环，性能远高于apply实现，适合大数据量场景。

示例数据构造 #

import pandas as pd
import numpy as np

# 构造示例DataFrame
df = pd.DataFrame([
    [1, 1, 0, 0, 0, 0, 0, 0, 0],
    [1, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0, 0, 0, 0],
    [1, 5, 5, 0, 0, 1, 0, 4, 6]
], columns=['c1','c2','c3','c4','c5','c6','c7','c8','c9'])

向量化实现代码 #

# 生成布尔掩码：标记所有值大于0的位置
mask = df.gt(0)
# 提取所有符合条件位置的行号、列号
row_idx, col_idx = np.where(mask)
# 按行号分组，收集每行对应的列索引
res = np.split(col_idx, np.unique(row_idx, return_index=True)[1])[1:]

结果验证 #

for item in res:
    print(item.tolist())

输出结果：

[0, 1]
[0]
[1]
[0, 1, 2, 5, 7, 8]

补充说明 #

如果需要和原DataFrame的行索引绑定，可以将结果转为Series：

res_series = pd.Series(res, index=df.index)

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