如何大幅提升Python多层嵌套循环的运算速度？

API 探路者

2026-10-4

多层嵌套循环优化方案

你这个计算逻辑可以完全通过向量化运算替代Python原生循环，能获得至少百倍以上的性能提升，以下是两种成熟实现方案：

方案1：使用numpy.einsum实现向量化计算（推荐）

爱因斯坦求和协定可以直接对应你这个多维度乘加逻辑，无需手动写任何循环，所有运算都在numpy底层C层面执行：

import numpy as np

# 先计算所有i,j组合的求和结果
B = np.einsum('klmn,ki,li,mj,nj->ij', x, r, r, r, r)
# 排除i=j的情况，按维度j求和得到最终values
values = B.sum(axis=1) - B[np.arange(10), np.arange(10)]

逻辑说明

einsum的下标字符串直接对应你公式里的维度匹配关系，自动完成所有维度的乘加操作
最后减去i==j位置的值，等价于原代码里的if i != j判断逻辑

方案2：使用numba JIT编译原生循环

如果后续逻辑会变得更复杂、不方便改写为向量化形式，可以直接用numba编译现有循环代码，改动量极小：

import numpy as np
from numba import njit

@njit  # 加这个装饰器即可，numba会自动编译函数为机器码
def calc_values(x, r):
    values = np.zeros(10, dtype=x.dtype)
    for i in range(10):
        val = 0
        for j in range(10):
            for k in range(10):
                for l in range(10):
                    for m in range(10):
                        for n in range(10):
                            if i != j:
                                val += x[k, l, m, n] * r[k, i] * r[l, i] * r[m, j] * r[n, j]
        values[i] = val
    return values

# 直接调用即可
values = calc_values(x, r)