Python处理大尺寸DataFrame生成匹配矩阵时如何降低时间复杂度

时间复杂度优化（直接降低数个量级） #

• 第一步砍掉无效对比：你的规则里只有start_elements和end_elements完全相同的行才需要进一步校验长度，所以先按这两列对数据分组，不同分组的行不需要做任何对比，直接默认值0。原逻辑是O(n²)复杂度，优化后复杂度变为所有分组内元素数量平方的总和，绝大多数场景下会直接降低3-5个数量级。
• 预计算常量：提前算好每行的长度length = end - start，不用循环中重复计算。不要在循环中反复调用iloc取数，提前把列值读入列表或numpy数组，单次读取即可。
• 分组内二次优化：如果单个分组内元素数量仍然很大，可以先将组内数据按length升序排序，对每个元素j，只需找长度大于等于0.2*length[j]的元素i（i<j），因为已经排序，可以用二分查找直接定位起始位置，无需遍历所有i，组内复杂度从O(k²)降到O(k logk)，k为分组内元素数量。
• 稀疏矩阵构建优化：不要用DOK矩阵逐点赋值，先收集所有需要填100的坐标对(i,j)，直接用scipy.sparse.coo_matrix批量构建，赋值速度比DOK快10倍以上。

内存优化补充建议 #

• 绝对不要调用todense()方法：百万级数据量下n*n的稠密矩阵会直接占用TB级内存，完全不可行，所有后续操作直接用稀疏矩阵完成即可。
• 用更小的数据类型：你的矩阵只有0和100两个值，可以用uint8类型存储（取值范围0-255），比默认的int32省75%内存。
• 可选择只存坐标对：如果不需要矩阵运算，直接存所有值为100的(i,j)对列表，内存占用比稀疏矩阵更低。

优化后代码示例 #

import pandas as pd
import numpy as np
from scipy.sparse import coo_matrix
import bisect

core_df = pd.DataFrame({'sno':[1,2,3,4],'start':[1,3,5,5],'end':[5,14,17,27],'start_elements':[('A', 'B'),('X','Y'),('B', 'C'),('B', 'C')],'end_elements':[('L', 'M'),('S', 'T'),('N', 'P'),('N', 'P')]})
# 预计算长度
core_df['length'] = core_df['end'] - core_df['start']
core_df_count = len(core_df)

# 收集所有需要填100的坐标
rows = []
cols = []

# 按start_elements和end_elements分组
for name, group in core_df.groupby(['start_elements', 'end_elements']):
    # 取出组内的原始索引和对应的长度，按长度排序
    idx_len = sorted(zip(group.index, group['length']), key=lambda x: x[1])
    idxs = [x[0] for x in idx_len]
    lens = [x[1] for x in idx_len]
    k = len(idxs)
    for j in range(1, k):
        cur_len = lens[j]
        min_len = 0.2 * cur_len
        # 二分查找第一个长度>=min_len的位置
        left = bisect.bisect_left(lens, min_len, 0, j)
        # 所有left到j-1的i都满足条件
        for i in range(left, j):
            rows.append(idxs[i])
            cols.append(idxs[j])

# 批量构建coo稀疏矩阵，值都是100，默认其余为0
data = [100] * len(rows)
score_matrix = coo_matrix((data, (rows, cols)), shape=(core_df_count, core_df_count), dtype=np.uint8)

# 只有数据量极小时才可以转成DataFrame，大数据量下注释掉下面两行
# Score_matrix_2D = pd.DataFrame(score_matrix.todense())
# print(Score_matrix_2D)

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