如何高效并行化Python中给数组赋值的简单for循环？

针对你提出的循环并行化需求，这里提供几种简洁且高效的实现方式，适配不同场景：

1. 标准库进程池（无需额外依赖） #

适合CPU密集型任务，通过concurrent.futures.ProcessPoolExecutor绕开GIL限制，语法简洁：

import numpy as np
from concurrent.futures import ProcessPoolExecutor

N = 100000

# 定义单元素计算逻辑
def compute_element(i):
    return i**2 + sum([j**(1/2) for j in range(10**4)])

if __name__ == "__main__":
    with ProcessPoolExecutor() as executor:
        # 并行处理所有i，结果按顺序返回
        results = executor.map(compute_element, range(N))
    # 转换为numpy数组
    a = np.fromiter(results, dtype=np.float64)

• 注意：Windows系统下必须把多进程代码放在if __name__ == "__main__"块内，避免进程启动异常
• 进程池默认使用CPU核心数作为进程数量，也可通过max_workers参数手动指定

2. Numba多线程（最接近Matlab/Julia的语法） #

如果是数值密集型任务，Numba是最优选择，语法和原循环几乎一致，且编译后效率极高：

首先安装Numba：pip install numba

import numpy as np
from numba import njit, prange

N = 100000
a = np.zeros(N)

# 启用并行编译
@njit(parallel=True)
def fill_array(a):
    # 用prange替代range，标记并行循环
    for i in prange(N):
        a[i] = i**2 + sum([j**(1/2) for j in range(10**4)])

fill_array(a)

• prange告诉Numba对该循环做并行优化，自动分配线程
• 函数会被编译为机器码，释放GIL，性能远超普通Python循环
• 支持原地修改数组，和原代码结构完全匹配，学习成本极低

3. 多进程池（底层实现） #

如果偏好更底层的多进程API，可使用multiprocessing.Pool：

import numpy as np
from multiprocessing import Pool

N = 100000

def compute_element(i):
    return i**2 + sum([j**(1/2) for j in range(10**4)])

if __name__ == "__main__":
    with Pool() as pool:
        results = pool.map(compute_element, range(N))
    a = np.array(results)

补充：向量化优化（非并行但更高效） #

如果你的实际任务逻辑可向量化，优先用Numpy向量化操作，速度通常比并行还快：

import numpy as np

N = 100000
# 提前预计算固定耗时部分，避免重复计算
fixed_sum = np.sum(np.sqrt(np.arange(10**4)))
# 直接生成数组，无循环
a = np.arange(N)**2 + fixed_sum

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