优化Python嵌套for循环速度与内存占用的技术求助

先纠正原代码的明显逻辑错误：循环变量n和m覆盖了外层定义的规模变量，导致循环完全不符合预期，修正后的基础代码如下：

n = 10000  # 可设置更大值
m = 10000  # 可设置更大值
x = []
y = []
for a in range(1, n+1):
    for b in range(1, m+1):
        x.append(a + b)
        y.append(a + b*2)

一、提速首选：NumPy向量化运算 #

Python原生循环的开销极大，用NumPy的向量化操作直接替代嵌套循环，底层靠C实现运算，速度能提升几十到上百倍：

import numpy as np

n = 10000
m = 10000

# 生成a和b的网格数组，避免循环
a_arr = np.repeat(np.arange(1, n+1), m)
b_arr = np.tile(np.arange(1, m+1), n)

# 向量化计算，一步到位
x = a_arr + b_arr
y = a_arr + b_arr * 2

如果数据规模大到NumPy数组也撑爆内存，就结合下面的分块/生成器方案。

二、低内存核心方案：生成器 #

如果不需要一次性把所有数据存在内存里，而是要逐个/批量处理结果，用生成器可以做到几乎零内存占用：

1. 逐元素生成 #

def generate_data(n, m):
    for a in range(1, n+1):
        for b in range(1, m+1):
            yield (a + b, a + b*2)

# 使用示例：逐个处理结果，比如写入文件
for x_val, y_val in generate_data(10000, 10000):
    # 这里写处理单个(x,y)的逻辑
    pass

2. 批量生成（适合批量写入/处理） #

def generate_batch(n, m, batch_size=10000):
    batch_x = []
    batch_y = []
    count = 0
    for a in range(1, n+1):
        for b in range(1, m+1):
            batch_x.append(a + b)
            batch_y.append(a + b*2)
            count += 1
            if count >= batch_size:
                yield (batch_x, batch_y)
                batch_x, batch_y, count = [], [], 0
    # 输出最后一批不足量的数据
    if batch_x:
        yield (batch_x, batch_y)

# 使用示例：批量处理
for batch_x, batch_y in generate_batch(10000, 10000):
    # 比如批量写入CSV、数据库
    pass

三、混合方案：分块处理+NumPy #

面对超大规模数据时，把数据分成多个小块，每个块用NumPy处理，既保证速度又严格控制内存占用：

import numpy as np

n = 10000
m = 10000
block_size = 1000  # 每个块处理的a的数量，可根据内存调整

for a_start in range(1, n+1, block_size):
    a_end = min(a_start + block_size - 1, n)
    # 生成当前块的a、b数组
    a_arr = np.repeat(np.arange(a_start, a_end+1), m)
    b_arr = np.tile(np.arange(1, m+1), a_end - a_start + 1)
    # 计算当前块的x、y
    block_x = a_arr + b_arr
    block_y = a_arr + b_arr * 2
    # 处理当前块数据，比如追加写入文件
    # np.savetxt('x_data.txt', block_x, mode='a')
    # np.savetxt('y_data.txt', block_y, mode='a')

四、原生循环的小优化（迫不得已时用） #

如果必须用原生Python循环，这两个小技巧能小幅提速：

• 预分配列表容量：避免列表动态扩容的开销

n = 10000
m = 10000
total = n * m
x = [0] * total
y = [0] * total
idx = 0
for a in range(1, n+1):
    for b in range(1, m+1):
        x[idx] = a + b
        y[idx] = a + b*2
        idx += 1

• 绑定局部变量：减少属性查找的开销

x = []
y = []
x_append = x.append
y_append = y.append
for a in range(1, n+1):
    for b in range(1, m+1):
        x_append(a + b)
        y_append(a + b*2)

内容的提问来源于stack exchange，提问作者Per Helge Semb