如何在Python中高效排序百万级数据集？求最优方案与内存优化技巧

1. 使用numpy（数值型数据集首选） #

numpy数组内存连续存储，排序逻辑由C实现，比原生sort()快数倍，完美适配百万级纯数值数据。

示例代码：

import numpy as np

# 生成百万级随机数值
data = np.random.randint(0, 1000000, size=1000000)
# 默认快速排序，可指定kind='mergesort'/'heapsort'
sorted_data = np.sort(data)

按自定义维度/字段排序时，结合argsort()更高效：

# 二维数组按第二列排序
data = np.random.rand(1000000, 2)
sorted_indices = np.argsort(data[:, 1])
sorted_data = data[sorted_indices]

2. 使用pandas（结构化数据集首选） #

针对带字段的记录（如字典列表、CSV数据），pandas的sort_values()依赖numpy优化，性能远超原生sorted()。

示例代码：

import pandas as pd

# 构造百万级结构化数据
data = pd.DataFrame({
    'id': range(1000000),
    'value': np.random.randint(0, 1000000, size=1000000),
    'category': np.random.choice(['A', 'B', 'C'], size=1000000)
})
# 按value列升序排序
sorted_data = data.sort_values(by='value')

超大CSV可分块读取排序，避免内存溢出：

chunk_size = 100000
sorted_chunks = []
# 分块读取并排序
for chunk in pd.read_csv('large_data.csv', chunksize=chunk_size):
    sorted_chunks.append(chunk.sort_values(by='value'))
# 合并所有排序后的块
final_sorted = pd.concat(sorted_chunks).sort_values(by='value')

3. heapq（Top-K场景优化） #

若无需全量排序，仅需前N个最大/最小元素，heapq的nlargest()/nsmallest()时间复杂度为O(n log k)，比全排序高效得多。

示例代码：

import heapq

# 百万级数据列表
data = [np.random.randint(0, 1000000) for _ in range(1000000)]
# 获取前100个最大元素
top_100 = heapq.nlargest(100, data)
# 获取前50个最小元素
bottom_50 = heapq.nsmallest(50, data)

自定义排序key的场景：

# 字典列表按value字段取前200个最大元素
data = [{'id': i, 'value': np.random.randint(0, 1000000)} for i in range(1000000)]
top_200 = heapq.nlargest(200, data, key=lambda x: x['value'])

4. 外部排序（超内存数据集） #

当数据量超过可用内存时，采用分治策略：分割为内存可承载的小块分别排序，写入临时文件后再归并。

简化版示例代码：

import os
import heapq

def split_and_sort(input_file, chunk_size=100000):
    chunk_num = 0
    with open(input_file, 'r') as f:
        while True:
            chunk = []
            for _ in range(chunk_size):
                line = f.readline()
                if not line:
                    break
                chunk.append(int(line.strip()))
            if not chunk:
                break
            chunk.sort()
            # 写入临时文件
            with open(f'temp_chunk_{chunk_num}.txt', 'w') as temp_f:
                temp_f.write('\n'.join(map(str, chunk)) + '\n')
            chunk_num += 1
    return chunk_num

def merge_sorted_chunks(chunk_num, output_file):
    temp_files = [open(f'temp_chunk_{i}.txt', 'r') for i in range(chunk_num)]
    heap = []
    # 初始化堆
    for i, f in enumerate(temp_files):
        line = f.readline()
        if line:
            heapq.heappush(heap, (int(line.strip()), i))
    
    with open(output_file, 'w') as out_f:
        while heap:
            val, idx = heapq.heappop(heap)
            out_f.write(f'{val}\n')
            # 读取对应文件下一行
            next_line = temp_files[idx].readline()
            if next_line:
                heapq.heappush(heap, (int(next_line.strip()), idx))
    
    # 清理临时文件
    for f in temp_files:
        f.close()
    for i in range(chunk_num):
        os.remove(f'temp_chunk_{i}.txt')

# 使用示例
split_and_sort('large_data.txt')
merge_sorted_chunks(10, 'sorted_large_data.txt')

• 用生成器代替列表加载数据：逐行生成数据，避免一次性加载全量数据到内存：

  def load_data(file_path):
      with open(file_path, 'r') as f:
          for line in f:
              yield int(line.strip())
  

• 选择内存高效的数据结构：

  • 用numpy数组替代Python列表：内存占用仅为列表的1/5左右（以int类型为例）。
  • 用pandas Categorical类型存储重复字符串：大幅降低类别字段的内存开销：

    data['category'] = pd.Categorical(data['category'])
    

• 分块处理数据：无论读取还是排序，每次仅处理一小部分数据，避免内存溢出。

• 减少排序key的开销：

  • 用operator.itemgetter代替lambda，性能更高：

    from operator import itemgetter
    sorted_data = sorted(data, key=itemgetter('value'))
    

  • 预计算排序key，避免排序时重复计算：

    data_with_key = [(x['value'], x) for x in data]
    data_with_key.sort()
    sorted_data = [x[1] for x in data_with_key]
    

• 避免不必要的拷贝：使用列表原地排序list.sort()（比sorted()少一次拷贝），numpy用np.ndarray.sort()实现原地排序。

内容的提问来源于stack exchange，提问作者I AM 1194