如何获取DataFrame指定列中含特定值的行号？

需求 #

获取指定列中包含特定列表值的行号，已知该值在列中唯一存在。

尝试代码 #

import pandas as pd
pos = pd.DataFrame(columns=['id', 'pred'])
pos.loc[1,'id'] = [4, 4, 4]
pos.loc[2,'id'] = [2, 3, 3]
pos.loc[3,'id'] = [1, 1, 2]
pos.loc[4,'id'] = [1, 2, 1]

print(pos)

print(pos[pos['id'] == [1, 1, 2]].index[0])

报错信息 #

ValueError: ('Lengths must match to compare', (6,), (3,))
中文翻译：值错误：比较时长度必须匹配，当前参与对比的长度为(6,)和(3,)

问题原因 #

直接用==对比DataFrame列和列表时，pandas会错误地将列中所有列表元素展开成一维序列，再与目标列表做逐元素匹配，导致两者长度不匹配，触发报错。

解决方案 #

要匹配列中存储的完整列表值，需逐行检查每个元素是否等于目标列表，推荐用apply方法实现：

import pandas as pd
pos = pd.DataFrame(columns=['id', 'pred'])
pos.loc[1,'id'] = [4, 4, 4]
pos.loc[2,'id'] = [2, 3, 3]
pos.loc[3,'id'] = [1, 1, 2]
pos.loc[4,'id'] = [1, 2, 1]

target_list = [1, 1, 2]
# 逐行判断列中的列表是否等于目标列表
matched_index = pos[pos['id'].apply(lambda x: x == target_list)].index[0]
print(matched_index)

运行后会输出目标值所在的行号3。

也可以用eq结合all的方式实现：

matched_index = pos[pos['id'].apply(lambda x: pd.Series(x).eq(target_list).all())].index[0]

这种方式将每个列表转为Series，逐元素对比后判断是否全部匹配，结果与前一种方法一致。

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