如何优化17万行DataFrame自定义索引算法的效率？

我有一个从CSV导入的DataFrame（含多列），共172033行。我写了一个自定义索引函数，用来筛选name属性相似度足够高的记录对，但算法效率极低——仅迭代10次就耗时约1分钟，全量运行时间完全无法接受。

当前自定义索引函数代码：

class CustomIndex(BaseIndexAlgorithm):
    def _link_index(self, df_a, df_b):
        indici1=[]
        indici2=[]
        for i in range(0, 173033):
            if(i%2 == 0):
                print(i) # 跟踪迭代进度
            for j in range(i, 173033):
                if(similar(df_a.loc[i, 'name'], df_a.loc[j, 'name'])>0.5):
                    indici1.append(i)
                    indici2.append(j)
        
        indici = [indici1, indici2]
        return pd.MultiIndex.from_arrays(indici, names=('first', 'second'))

相似度函数：

from difflib import SequenceMatcher

def similar(a, b):
    return SequenceMatcher(None, a, b).ratio()

需求：生成一个MultiIndex对象，包含所有相似度>0.5的记录对（包括自身匹配，如(0,0)，以及相似对如(0,3)、(1,4)这类）。

DataFrame示例：

name
0  Amazon
1  Walmart
2  Apple
3  Amazon.com
4  Walmart Inc.

你的代码核心问题是O(n²)的双重循环，加上SequenceMatcher本身的计算开销，对于17万行数据来说完全不可行。以下是从基础到进阶的优化方法：

1. 先做文本预处理，减少无效计算 #

先清洗name列，统一格式、移除干扰后缀，能大幅降低相似度计算的复杂度，同时提升匹配准确性：

def clean_text(s):
    s = s.lower().strip()
    # 移除常见干扰后缀
    suffixes = ['.com', ' inc.', ' llc.', ' corp.', ' limited']
    for suf in suffixes:
        if s.endswith(suf):
            s = s[:-len(suf)].strip()
    return s

# 应用到DataFrame
df['cleaned_name'] = df['name'].apply(clean_text)

2. 用批量矩阵运算替代双重循环 #

放弃逐行逐对的循环，改用底层优化的矩阵运算，速度能提升几个数量级：

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import pairwise_distances
import numpy as np

# 把清洗后的文本转成TF-IDF向量
vectorizer = TfidfVectorizer(ngram_range=(1,2))
tfidf_matrix = vectorizer.fit_transform(df['cleaned_name'])

# 计算余弦相似度（1 - 余弦距离）
similarity_matrix = 1 - pairwise_distances(tfidf_matrix, metric='cosine')

# 筛选相似度>0.5且i<=j的记录对（只保留上三角+对角线）
rows, cols = np.where((similarity_matrix > 0.5) & (np.triu(np.ones(similarity_matrix.shape), k=0) == 1))

# 生成目标MultiIndex
result_index = pd.MultiIndex.from_arrays([rows, cols], names=('first', 'second'))

3. 替换低效的相似度计算函数 #

SequenceMatcher效率偏低，换成更快的实现：

• 用rapidfuzz（fuzzywuzzy的C语言实现，速度提升10-100倍）
  安装依赖：

pip install rapidfuzz

替换后的相似度函数：

from rapidfuzz import fuzz

def similar(a, b):
    return fuzz.ratio(a, b) / 100  # 转成0-1的比例值

4. 用分组预筛选减少计算量 #

如果不需要绝对精确的所有相似对，可以先通过n-gram倒排索引把相似文本分组，只在组内计算相似度，进一步减少计算量：

from collections import defaultdict

# 生成n-gram倒排索引，把含相同n-gram的文本归为候选组
def get_ngrams(s, n=2):
    return set([s[i:i+n] for i in range(len(s)-n+1)])

ngram_index = defaultdict(list)
for idx, name in df['cleaned_name'].items():
    for ngram in get_ngrams(name):
        ngram_index[ngram].append(idx)

# 生成去重的候选对（仅保留i<=j）
candidate_pairs = set()
for idx in df.index:
    name = df.loc[idx, 'cleaned_name']
    candidates = set()
    for ngram in get_ngrams(name):
        candidates.update(ngram_index[ngram])
    for candidate in candidates:
        if idx <= candidate:
            candidate_pairs.add((idx, candidate))

# 对候选对计算精确相似度
indici1, indici2 = [], []
for i, j in candidate_pairs:
    if fuzz.ratio(df.loc[i, 'cleaned_name'], df.loc[j, 'cleaned_name']) > 50:
        indici1.append(i)
        indici2.append(j)

result_index = pd.MultiIndex.from_arrays([indici1, indici2], names=('first', 'second'))

5. 细节优化 #

• 移除循环内的print(i)，频繁IO会拖慢速度，换成tqdm显示进度条：

  from tqdm import tqdm
  for i in tqdm(range(len(df))):
      ...
  

• 提前把cleaned_name转成列表，避免循环中重复调用df.loc：

  names_list = df['cleaned_name'].tolist()
  # 循环内直接用names_list[i]、names_list[j]
  

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