如何在Pandas中高效实现城市变体的部分字符串匹配？

我有一个含约5万条城市名称变体的DataFrame，示例如下：

另有一个包含200万条变体的倒排索引表，示例如下：

我需要计算第一个DataFrame中每个变体在倒排索引中至少部分出现的总计数（例如new dilli对应倒排索引里的3条记录，总计数是5+4+3=12）。

我写了以下代码，能正常运行但速度极慢，求优化方案：

def if_contains(pattern,df,col,sum_col):
    return df[df[col].str.contains(pattern)][sum_col].sum()

inputs["count_of_vns_with_pattern"] = inputs.apply(
    lambda x: if_contains(
        x["variation"], inv_index, "city_variations", "No.of times present in inv index"
    ),
    axis=1,
)

1. 批量正则匹配（核心优化，大幅提速） #

原代码的时间复杂度是O(5万*200万)，属于暴力遍历。可以把所有待匹配的变体合并成一个正则表达式，一次性完成倒排索引的匹配，再反向汇总结果：

import re
from collections import defaultdict

# 转义变体中的特殊字符，避免正则语法冲突
unique_variations = inputs["variation"].unique()
escaped_patterns = [re.escape(var) for var in unique_variations]
# 合并为一个正则，匹配任意变体
combined_re = re.compile("|".join(escaped_patterns))

# 给倒排索引的每条记录标记匹配到的所有变体
inv_index["matched_vars"] = inv_index["city_variations"].str.findall(combined_re)

# 遍历倒排索引，汇总每个变体的总计数
count_dict = defaultdict(int)
for _, row in inv_index.iterrows():
    for var in row["matched_vars"]:
        count_dict[var] += row["计数"]

# 映射回原DataFrame
inputs["count_of_vns_with_pattern"] = inputs["variation"].map(count_dict).fillna(0)

该方法时间复杂度为O(200万 + 5万)，比原方法效率提升几个数量级。

2. 预编译正则（小幅提速） #

如果必须保留逐行匹配逻辑，至少预编译所有变体的正则表达式，避免每次调用str.contains重复编译：

import re

# 预编译所有变体的正则
pattern_map = {var: re.compile(re.escape(var)) for var in inputs["variation"].unique()}

def calc_match_sum(pattern, df, col, sum_col):
    return df[df[col].str.contains(pattern)][sum_col].sum()

# 用map代替apply，减少lambda的额外开销
inputs["count_of_vns_with_pattern"] = inputs["variation"].map(
    lambda x: calc_match_sum(pattern_map[x], inv_index, "city_variations", "计数")
)

3. 全文索引工具（超大数据量场景） #

如果200万条数据仍有性能压力，可以用专门的全文索引库Whoosh建立索引，实现快速模糊匹配：

from whoosh.index import create_in
from whoosh.fields import Schema, TEXT, NUMERIC
from whoosh.qparser import QueryParser
import os

# 创建索引目录和结构
if not os.path.exists("city_index"):
    os.mkdir("city_index")
schema = Schema(city_var=TEXT(stored=True), count=NUMERIC(stored=True))
ix = create_in("city_index", schema)

# 写入倒排索引数据
writer = ix.writer()
for _, row in inv_index.iterrows():
    writer.add_document(city_var=row["city_variations"], count=row["计数"])
writer.commit()

# 查询每个变体的匹配计数
count_dict = {}
with ix.searcher() as searcher:
    parser = QueryParser("city_var", ix.schema)
    for var in inputs["variation"].unique():
        # 用通配符匹配包含当前变体的记录
        query = parser.parse(f"*{var}*")
        results = searcher.search(query)
        total = sum(res["count"] for res in results)
        count_dict[var] = total

inputs["count_of_vns_with_pattern"] = inputs["variation"].map(count_dict).fillna(0)

关键优化逻辑 #

• 避免逐行遍历+全量扫描：原方法的核心问题是重复对200万条数据做全量匹配，优化后将复杂度从O(N*M)降到O(N+M)
• 正则预编译：减少重复编译正则表达式的额外开销
• 批量处理：优先处理倒排索引，再将结果映射回原DataFrame，减少重复计算

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