如何快速遍历不等长列查找存在值？Pandas两DataFrame匹配效率优化

原代码采用三层嵌套遍历逻辑，时间复杂度为O(df1行数 × df1单列表元素数 × df2行数)，3.5万行df1 + 7.6万行df2的场景下会产生超260亿次判断操作，因此运行极慢。

核心优化思路 #

• 先对df2构建倒排索引：将每个token作为key，value为所有包含该token的col2列表的集合，构建过程仅需遍历df2所有子元素1次
• 直接遍历df1的所有token，通过倒排索引直接命中对应结果，无需二次遍历df2

优化后代码 #

import pandas as pd
from collections import defaultdict

def query_bounce_checker(dataset_clicked, dataset_loaded, col1, col2):
    # 第一步：构建倒排索引
    token_index = defaultdict(list)
    for lst in dataset_loaded[col2]:
        for token in lst:
            token_index[token].append(lst)
    # 第二步：查询匹配结果
    query_bounce = []
    for sub_lst in dataset_clicked[col1]:
        for token in sub_lst:
            if token in token_index:
                query_bounce.extend(token_index[token])
    return query_bounce

# 示例数据测试，结果和原代码完全一致
post_token_list = [['wXrL3TbK'], ['wXmTQKw1'], ['wXvnlWej'], ['wXvXBjKp']]
tokens_list = [['wXv3qoPQ', 'wXvT7ylu', 'wXvnIJuH', 'wXvXH7vy', 'wXvDXSS1', 'wXvjVE1F', 'wXvPV6z1', 'wXvHF1uw',
                'wXvH1q03', 'wXvnTlcr', 'wXvDEG9U', 'wXLfZtO6', 'wXvLDDDl', 'wXvHTgjk', 'wXvHDDr8', 'wXvPBLbu',
                'wXvvxXHI', 'wXvPBFge', 'wXvLxSii', 'wXvDhk2h', 'wXv3Alan', 'wXvvQuKy', 'wXvvQ6LO', 'wXpHNjw9'],
               ['wXYr2lVk', 'wXXj7iDP', 'wXXXIsQr', 'wXQbXKz6', 'wXN3tMp1', 'wXMfZV5N', 'wXvnlWej', 'wXSDyEaW',
                'wXQ7mM78', 'wXMPvojh', 'wXMjo-8G', 'wXLfZtO6', 'wXN3tMp1'],
               ['wXr_jZmX', 'wXr7D0AM', 'wXrzjhxL', 'wXrfjQNe', 'wXrnihqT', 'wXrjyqm5', 'wXr3CD4h', 'wXrnSZsy',
                'wXrTieP7', 'wXLfZtO6', 'wXgHVwkc', 'wXdvewsV', 'wXrfxZeg', 'wXrLB7Zo', 'wXprtX71', 'wXrHhjtO',
                'wXrzwKBt', 'wXqz-RlY', 'wXq_fp7F', 'wXq7Po7n', 'wXq7fC73', 'wXqzvRSW', 'wXqf_PQ3', 'wXML2vCd'],
               ['wXv3aQrv', 'wXvn6ONM', 'wXvfaG0M', 'wXvf6LIr', 'wXvjJBg_', 'wXvL6M-0', 'wXv7p2cd', 'wXv3poSs',
                'wXvz5kUz', 'wXvrZz0_', 'wXv_YVCb', 'wXLfZtO6', 'wXvX5Hgi', 'wXvz3Ptg', 'wXvHJUU-', 'wXvr4fB7',
                'wXvnlWej', 'wXv_YUrK', 'wXv7Id05', 'wXv7IYOV', 'wXvfYfLo', 'wXv7Y3AV', 'wXvT4_pE', 'wXvPovRt'],
               ['wXoDui-2', 'wXoT9yTg', 'wXmTQKw1', 'wXormLxu', 'wXMX-NNQ', 'wXo7kUfB', 'wXon0rt_', 'wXozT-3V',
                'wXnvYjEc', 'wXnTn9D6', 'wXnLH7Cz', 'wXn_2HV_', 'wXnPGou9', 'wXnPVSNo', 'wXuG0sl3', 'wXnjAs7X',
                'wXm38mLv', 'wXmnj5Oh', 'wXmfjQ2h', 'wXm_wXuD', 'wXlPOUmy', 'wXcfHkmx', 'wXQ_62cx', 'wXUD3qyx']]

df1 = pd.DataFrame({"col1": post_token_list})
df2 = pd.DataFrame({"col2": tokens_list})

res = query_bounce_checker(df1, df2, "col1", "col2")

性能说明 #

• 假设df2每行平均有24个token，7.6万行仅需遍历约182万次即可完成索引构建
• df1 3.5万行仅需遍历3.5万次即可完成所有查询
• 全量数据场景下运行耗时可从小时级降到毫秒/秒级，结果和原代码完全一致

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