如何用NumPy优化嵌套四层for循环的复数累加计算？

你的四层循环本质是计算双重求和的乘积，先从数学上拆解问题，再用NumPy的向量化操作彻底消除循环。

1. 数学拆解（核心优化点） #

原式中，累加项可以拆分为两个独立求和的乘积：

a[i,j] = sum_{k=0}^{N-1} sum_{l=0}^{N-1} exp(2πi/N (i*k + j*l))
       = [sum_{k=0}^{N-1} exp(2πi * i*k / N)] × [sum_{l=0}^{N-1} exp(2πi * j*l / N)]

因为exp(a+b) = exp(a)×exp(b)，双重求和可以拆成两个一维求和的外积，复杂度从O(N⁴)直接降到O(N²)。

2. NumPy向量化实现 #

基于上述拆解，用广播机制完成计算，完全消除显式循环：

import numpy as np

N = 5
# 生成索引数组
i_indices = np.arange(N)
j_indices = np.arange(N)

# 计算每个i对应的k求和项
sum_k = np.exp(2j * np.pi / N * i_indices[:, None] * i_indices).sum(axis=1)
# 计算每个j对应的l求和项
sum_l = np.exp(2j * np.pi / N * j_indices[:, None] * j_indices).sum(axis=1)

# 外积得到最终矩阵
a = sum_k[:, None] * sum_l[None, :]

这里i_indices[:, None]将一维数组转为列向量，和行向量i_indices广播成N×N矩阵，对应所有i和k的组合，求和后得到每个i的结果；同理处理sum_l，最后通过广播相乘得到完整的N×N矩阵。

3. 极致优化：利用离散傅里叶正交性 #

观察求和项的数学性质：
sum_{k=0}^{N-1} exp(2πi * m*k / N) 是离散傅里叶级数的正交基求和，结果为：

• 当m是N的倍数时，结果为N；
• 否则结果为0。

在你的代码中，m=i（0≤i<N），因此只有当i=0时sum_k=N，其余i的sum_k=0；同理sum_l只有j=0时为N。因此最终矩阵仅(0,0)位置为N×N=N²，其余全为0，可以直接构造：

import numpy as np

N = 5
a = np.zeros((N, N), dtype=np.complex128)
a[0, 0] = N * N

这个版本效率最高，完全跳过了求和计算。

验证正确性 #

你可以对比原始循环代码和上述优化代码的输出，结果完全一致，但优化后的代码在N较大时（比如N=100），速度会提升几个数量级。

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