Python百万级数据集嵌套循环性能优化方案咨询

你的代码核心问题是嵌套循环带来的O(n²)时间复杂度，当数据集达到百万行时，这种方法完全不可行。以下是几种高效的优化方案：

1. 字典预分组（最推荐，时间复杂度O(n)） #

思路：先按对比条件（第一列的值）将所有行的索引分组，再在每个组内生成符合要求的配对。这样避免了逐行对比的嵌套循环，把时间复杂度从O(n²)降到线性级别。

示例代码：

import csv
from collections import defaultdict
from itertools import combinations

file_path = 'data.csv'

# 按第一列的值分组存储索引
index_groups = defaultdict(list)
with open(file_path, 'r') as file:
    reader = csv.reader(file)
    for idx, row in enumerate(reader):
        index_groups[row[0]].append(idx)

matching_pairs = []
# 遍历每个分组，生成所有i < j的索引配对
for indices in index_groups.values():
    # 和原代码逻辑一致，仅存储配对中的i
    for i, j in combinations(indices, 2):
        matching_pairs.append(i)
    # 若需存储完整(i,j)配对，替换为：
    # matching_pairs.extend(combinations(indices, 2))

output_file = 'matching_pairs.txt'
with open(output_file, 'w') as file:
    for pair in matching_pairs:
        file.write(f'{pair}\n')

2. 使用Pandas处理（适合数据分析场景） #

Pandas的groupby操作基于C实现，处理大数据集的效率远高于纯Python循环，适合已经使用数据分析技术栈的场景。

示例代码：

import pandas as pd
from itertools import combinations

df = pd.read_csv('data.csv', header=None)
matching_pairs = []

# 按第一列分组，遍历每个组的索引
for _, group in df.groupby(0):
    indices = group.index.tolist()
    for i, j in combinations(indices, 2):
        matching_pairs.append(i)

# 写入文件
pd.Series(matching_pairs).to_csv('matching_pairs.txt', index=False, header=False)

3. Numba加速循环（保留类似循环逻辑时使用） #

如果你的对比逻辑比示例更复杂、无法用分组替代，可以用Numba将Python循环编译为机器码，大幅提升执行速度。注意Numba对纯Python语法支持较好，避免使用复杂的数据结构。

示例代码：

import csv
from numba import njit

file_path = 'data.csv'

data = []
with open(file_path, 'r') as file:
    reader = csv.reader(file)
    data = list(reader)

# 用Numba装饰器编译函数
@njit
def find_matching_pairs(data):
    n = len(data)
    matching_pairs = []
    for i in range(n):
        key_i = data[i][0]
        for j in range(i + 1, n):
            if data[j][0] == key_i:
                matching_pairs.append(i)
    return matching_pairs

matching_pairs = find_matching_pairs(data)

output_file = 'matching_pairs.txt'
with open(output_file, 'w') as file:
    for pair in matching_pairs:
        file.write(f'{pair}\n')

4. 内存优化（针对超大规模数据集） #

如果数据集大到无法一次性加载到内存，可以分块处理：

• 使用csv.reader逐行处理并分组，避免一次性存入列表
• 使用Pandas的chunksize参数分块读取，再合并分组结果

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