如何加速Python中生成矩阵z的4层嵌套循环运行速度？

• 核心思路：利用Numpy的向量化运算替代Python循环
  Python原生for循环在处理大规模数据时效率极低，而Numpy基于C实现的向量化操作能大幅提升计算速度。我们可以拆解原公式，借助Numpy的广播机制直接生成目标矩阵，完全规避嵌套循环。

• 具体实现步骤

  1. 预计算x相关固定项：提前算出x[i]**alpha + (1 - delta) * x[i]的数组，减少重复计算
  2. 利用广播构造四维数组，对应原循环中i、j、a、b四个维度
  3. 将四维数组reshape为目标形状(T*T, T*T)的矩阵z

• 优化后的代码

import numpy as np

T        =  50
beta     =  0.98 
alpha    =  0.3
delta    =  0.1

xss      =  ((1 / beta - (1 - delta)) / alpha)**(1 / (alpha - 1))
yss      =  xss**alpha

x        =  np.linspace(0.8 * xss, 1.4 * xss, T,dtype='float32')
y        =  np.linspace(yss, 1.2 * yss, T,dtype='float32')

# 预计算x的固定项：x[i]^alpha + (1-delta)*x[i]
term_x = x ** alpha + (1 - delta) * x

# 利用广播构造四维数组，对应i,j,a,b四个维度
z_4d = term_x[:, None, None, None] + y[None, :, None, None] - x[None, None, :, None] - y[None, None, None, :]

# 转换为目标形状(T*T, T*T)
z = z_4d.reshape(T*T, T*T)

• 另一种更直观的拆分方式
  可将原公式拆分为(term_x[i] + y[j])和(-x[a] - y[b])两部分，分别构造二维矩阵后通过广播相加：

# 构造(T,T)的矩阵A，对应term_x[i] + y[j]
A = term_x[:, None] + y[None, :]
# 构造(T,T)的矩阵B，对应 -x[a] - y[b]
B = -x[:, None] - y[None, :]

# 广播相加得到四维数组，再reshape
z_4d = A[:, :, None, None] + B[None, None, :, :]
z = z_4d.reshape(T*T, T*T)

• 效果对比
  当T=50时，原循环需执行50^4=6,250,000次迭代，优化后的向量化操作几乎瞬间完成；当T更大（如T=100），原循环会极其缓慢，向量化方案的性能优势会更显著。

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