如何快速基于多列列表值筛选Pandas DataFrame的索引？

问题背景 #

现有包含col1和col2两列的Pandas DataFrame，需根据两个列表l1与l2的对应组合值筛选数据，提取匹配行的索引。当前循环遍历的实现方式在处理大型数据集时速度过慢，需要更高效的替代方案。

示例数据与低效实现 #

import pandas as pd
d = {'col1': [11, 20,90,80,30], 'col2': [30, 40,50,60,90]}
df = pd.DataFrame(data=d)
print(df)
#    col1  col2
# 0    11    30
# 1    20    40
# 2    90    50
# 3    80    60
# 4    30    90

l1=[11,90,30]
l2=[30,50,90]
final_result=[]
for i,j in zip(l1,l2):
    res=df[(df['col1']==i) & (df['col2']==j)]
    final_result.append(res.index[0])
print(final_result)  # 输出 [0, 2, 4]

高效实现方案 #

方案1：通过pd.merge合并匹配 #

将l1和l2构造为临时DataFrame，利用Pandas的合并操作完成批量匹配，同时保留原DataFrame的索引。这种向量化操作完全规避了循环，处理大型数据时效率提升显著。

import pandas as pd

d = {'col1': [11, 20,90,80,30], 'col2': [30, 40,50,60,90]}
df = pd.DataFrame(data=d)
l1=[11,90,30]
l2=[30,50,90]

# 构造匹配用的临时DataFrame
match_df = pd.DataFrame({'col1': l1, 'col2': l2})
# 合并原DataFrame并保留原索引
merged = pd.merge(match_df, df.reset_index(), on=['col1', 'col2'], how='left')
# 提取结果索引列表
final_result = merged['index'].tolist()
print(final_result)  # 输出 [0, 2, 4]

方案2：利用元组序列结合isin匹配 #

将原DataFrame的col1和col2转换为元组序列，直接用isin匹配l1与l2组成的元组列表，快速筛选出目标行并提取索引。

import pandas as pd

d = {'col1': [11, 20,90,80,30], 'col2': [30, 40,50,60,90]}
df = pd.DataFrame(data=d)
l1=[11,90,30]
l2=[30,50,90]

# 构造匹配元组列表
match_tuples = list(zip(l1, l2))
# 筛选匹配行并提取索引
final_result = df[df[['col1', 'col2']].apply(tuple, axis=1).isin(match_tuples)].index.tolist()
print(final_result)  # 输出 [0, 2, 4]

方案3：设置复合索引后直接匹配 #

将col1和col2设为复合索引，直接通过元组列表索引目标行，再提取原DataFrame的索引。该方法在数据量较大时性能表现最优。

import pandas as pd

d = {'col1': [11, 20,90,80,30], 'col2': [30, 40,50,60,90]}
df = pd.DataFrame(data=d)
l1=[11,90,30]
l2=[30,50,90]

# 设置复合索引，保留原索引
df_multi = df.set_index(['col1', 'col2'])
# 构造匹配元组
match_tuples = list(zip(l1, l2))
# 提取对应的原索引
final_result = df_multi.loc[match_tuples].index.get_level_values(0).tolist()
print(final_result)  # 输出 [0, 2, 4]

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