如何在pandas DataFrame中查找包含特定值的列并返回每行对应列索引

解决方案 #

可以直接用pandas等值判断配合numpy的向量化操作实现，效率远高于逐行遍历：

1. 首先构造测试DataFrame：

import pandas as pd
import numpy as np

test = pd.DataFrame({"id":[1,2,3,4],
                     "name_1":["peter","bobby","alex","chris"],
                     "name_1_flag":["real","fake","fake","real"],
                     "name_2":["hector","abi","henrik","miko"],
                     "name_2_flag":["fake","real","fake","fake"],
                     "name_3":["hans","khan","will","than"],
                     "name_3_flag":["fake","fake","real","fake"]})

2. 生成要求的series结果：

# 生成1-based的列索引series
res = pd.Series(np.argmax(test.eq('real').to_numpy(), axis=1) + 1, name='col_index')

运行后res的输出和要求完全一致：

0    3
1    5
2    7
3    3
Name: col_index, dtype: int64

如果需要获取所有real值对应的行、列元组，可以用以下代码：

# 行列均为0-based的元组列表
row_col_tuples = list(zip(*np.where(test.eq('real'))))
# 输出：[(0, 2), (1, 4), (2, 6), (3, 2)]

# 列索引为1-based的元组列表
row_col_tuples_1based = [(r, c+1) for r,c in zip(*np.where(test.eq('real')))]
# 输出：[(0, 3), (1, 5), (2, 7), (3, 3)]

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