如何提升Python中依赖行列前置数据的双重for循环执行速度？

针对你的问题，分两种场景给出优化方案：

一、当前计算仅依赖已知矩阵（原代码场景） #

你的原代码中，result[n,j]的计算仅依赖预初始化好的matrix，不依赖result自身的其他值，这种情况可以直接用Numpy向量化操作完全替代双重循环，速度提升非常明显：

import numpy as np

num_rows, num_cols = 100, 5

matrix = np.zeros((num_rows, num_cols))
matrix[0, :] = np.random.rand(num_cols)
matrix[:, 0] = np.random.rand(num_rows)

coeff1 = np.random.rand(num_rows)
coeff2 = np.random.rand(num_rows)
coeff3 = np.random.rand(num_rows)

result = np.zeros_like(matrix)
# 复制初始行和列
result[0, :] = matrix[0, :]
result[:, 0] = matrix[:, 0]

# 向量化计算核心部分
# term1: coeff1[n] * matrix[n-1,j-1]，对应切片扩展维度后广播计算
term1 = coeff1[1:, np.newaxis] * matrix[:-1, :-1]
# term2: coeff2[n] * matrix[n,j-1]
term2 = coeff2[1:, np.newaxis] * matrix[1:, :-1]
# term3: coeff3[n] * matrix[n-1,j]
term3 = coeff3[1:, np.newaxis] * matrix[:-1, 1:]

# 赋值到result对应位置
result[1:, 1:] = term1 + term2 + term3

这里用Numpy的切片+广播操作一次性完成所有元素的计算，彻底避免Python层面的循环开销，大矩阵场景下速度提升会非常显著。

二、计算依赖自身矩阵的前一行/列（动态规划场景） #

如果你的实际场景是result[n,j]依赖result[n-1,j]（前一行当前列）、result[n,j-1]（当前行前一列）这类自身已计算的值，比如动态规划问题，这时候无法完全向量化，但可以用以下方法优化：

1. 使用Numba JIT编译 #

Numba可以将Python函数编译为机器码，大幅加速循环。只需要给循环函数加上装饰器即可：

import numpy as np
from numba import jit

num_rows, num_cols = 100, 5

matrix = np.zeros((num_rows, num_cols))
matrix[0, :] = np.random.rand(num_cols)
matrix[:, 0] = np.random.rand(num_rows)

coeff1 = np.random.rand(num_rows)
coeff2 = np.random.rand(num_rows)
coeff3 = np.random.rand(num_rows)

result = np.zeros_like(matrix)
result[0, :] = matrix[0, :]
result[:, 0] = matrix[:, 0]

# 用numba jit编译循环函数，nopython模式生成纯机器码
@jit(nopython=True)
def compute_result(result, coeff1, coeff2, coeff3):
    for j in range(1, num_cols):
        for n in range(1, num_rows):
            # 这里term3依赖result[n-1,j]，符合你的问题描述
            term1 = coeff1[n] * result[n-1, j-1]
            term2 = coeff2[n] * result[n, j-1]
            term3 = coeff3[n] * result[n-1, j]
            result[n, j] = term1 + term2 + term3

compute_result(result, coeff1, coeff2, coeff3)

nopython=True会强制Numba生成纯机器码，避免Python对象的开销，这类循环的加速效果通常能达到几十甚至上百倍。

2. 手动循环优化（减少属性访问） #

如果不能使用第三方库，可以手动优化循环内的操作，比如提前把数组转为局部变量，减少循环内的属性查找：

# 提前提取局部变量，减少循环内的全局/属性访问
res = result
c1 = coeff1
c2 = coeff2
c3 = coeff3

for j in range(1, num_cols):
    # 提前获取当前列的前一列数据，避免重复索引
    prev_col = res[:, j-1]
    for n in range(1, num_rows):
        term1 = c1[n] * res[n-1, j-1]
        term2 = c2[n] * prev_col[n]
        term3 = c3[n] * res[n-1, j]
        res[n, j] = term1 + term2 + term3

这种方法的提升幅度有限，但比原始循环要快一些，适合不能引入额外依赖的场景。

内容的提问来源于stack exchange，提问作者Vinicius B. de S. Moreira