如何消除多环求和计算冗余，提升大矩阵运算效率

问题核心 #

当前代码每次调用sum_of_values都会完整遍历整个矩阵，重复计算每个元素的U值（U = x² + x + y² + y + 1），当需要对数百甚至数千个d值（如0.5,1.5,...,999.5）计算时，冗余计算量线性增长，大矩阵下效率极低。

优化思路 #

核心是彻底消除重复遍历矩阵和重复计算U值的操作：

1. 预先一次性计算所有元素的U值，存储为矩阵；
2. 基于预计算的U矩阵，通过向量化操作或分组统计，一次性得到所有d对应的区间和。

高效实现方案 #

方案一：预计算U矩阵 + 循环掩码求和 #

利用NumPy向量化操作替代Python循环，大幅提升计算速度：

import numpy as np
from tqdm import tqdm

def calc_values_fast(matrix, n):
    # 转换为NumPy数组启用向量化计算
    matrix_np = np.array(matrix)
    # 生成对应原代码的x行、y列坐标网格
    x = np.arange(n)
    y = np.arange(n)
    X, Y = np.meshgrid(x, y, indexing='ij')
    
    # 一次性计算所有元素的U值
    U = X**2 + X + Y**2 + Y + 1
    
    scores = []
    # 遍历所有d值，用掩码筛选区间内元素求和
    for d in [i + 0.5 for i in range(n)]:
        lower_bound = (d ** 2) * 2
        upper_bound = ((d + 1) ** 2) * 2
        mask = (U > lower_bound) & (U < upper_bound)
        scores.append(matrix_np[mask].sum())
    return scores

方案二：预计算U矩阵 + 分组统计（更高效） #

通过数学变换直接映射每个元素到对应d区间，用bincount一次性完成所有区间求和，彻底消除循环：

import numpy as np

def calc_values_faster(matrix, n):
    matrix_np = np.array(matrix)
    x = np.arange(n)
    y = np.arange(n)
    X, Y = np.meshgrid(x, y, indexing='ij')
    U = X**2 + X + Y**2 + Y + 1
    
    # 数学变换：将U值映射到对应的d索引（d = i + 0.5，i为0到n-1）
    sqrt_U_half = np.sqrt(U / 2)
    d_indices = np.floor(sqrt_U_half - 0.5).astype(int)
    
    # 过滤超出有效范围的索引
    valid_mask = (d_indices >= 0) & (d_indices < n)
    
    # 分组求和，minlength保证结果长度为n
    sum_scores = np.bincount(
        d_indices[valid_mask],
        weights=matrix_np[valid_mask],
        minlength=n
    )
    return sum_scores.tolist()

性能对比（400×400矩阵） #

• 原代码：约26秒
• 方案一（循环掩码求和）：约149ms
• 方案二（分组统计）：约119ms

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