如何高效遍历矩阵对角线并赋值？大矩阵运行效率优化

我有多个2000×2000的矩阵，需要遍历每个矩阵的上/下对角线，提取元素并计算每条对角线元素的均值，随后将小于均值的元素置为0。此前尝试用for循环实现，但矩阵维度过大导致运行耗时过长，请问如何有效降低运行时间？

已尝试的代码实现 #

import numpy as np

def get_nth_diag_indices(mat, offset):
    rows, cols_orig = np.diag_indices_from(mat)
    cols = cols_orig.copy()
    if offset > 0:
        cols += offset
        rows = rows[:-offset]
        cols = cols[:-offset]
    return rows, cols

def normalize_along_diagonal_from_numpy(d, max_bin_distance, trim=0.01):
    for offset in range(1, max_bin_distance + 1):
        r, c = get_nth_diag_indices(d, offset)
        vals_orig = d[r, c].tolist()
        vals = vals_orig.copy()
        vals = list(filter(lambda num: num != 0, vals))
        vals.sort()
        vals.reverse()
        trim_index = round(trim * len(vals)) - 1
        if trim_index < 0:
            trim_index = 0
        remaining = vals[(trim_index):]
        if len(remaining) == 0:
            d[r, c] = [0] * len(vals_orig)
        else:
            mu = np.mean(remaining)
            sd = np.std(remaining)
            if sd < 1e-6:
                d[r, c] = [0] * len(vals_orig)
            else:
                d[r, c] = (d[r, c] - mu) / sd
                np.where(d[r, c] >= mu + sd, d[r, c], 0)
    ### replicate
    mat = d + d.T - np.diag(np.diag(d))
    return mat

优化方案 #

1. 用Numpy原生函数提取对角线，替代手动索引计算 #

Numpy的np.diag可以直接按偏移量提取对角线元素，np.fill_diagonal可直接修改对角线，比手动计算索引高效得多：

# 提取offset>0的上对角线元素
diag_vals = np.diag(mat, offset=offset)
# 将新值赋值回上对角线
np.fill_diagonal(mat[:, offset:], new_vals)

2. 全程用Numpy数组操作，避免Python列表转换 #

原代码中tolist()、filter、Python列表排序等操作在处理大数组时效率极低，全部替换为Numpy向量操作：

• 过滤0元素：non_zero = diag_vals[diag_vals != 0]
• 降序排序：sorted_vals = np.sort(non_zero)[::-1]
• 数值计算全程基于Numpy数组，减少Python层面的循环和判断开销

3. 修复原代码逻辑漏洞 #

原代码中np.where语句未将结果赋值回原矩阵，导致置0操作无效，需修改为：

d[r, c] = np.where(d[r, c] >= mu + sd, d[r, c], 0)

4. 减少重复计算，利用对称特性 #

只需要处理上对角线，最后通过矩阵转置复制结果到下对角线，无需单独遍历下对角线，直接减少一半循环次数。

优化后的示例代码 #

import numpy as np

def optimize_diagonal_processing(d, max_bin_distance, trim=0.01):
    mat = d.copy()
    n = mat.shape[0]
    
    for offset in range(1, max_bin_distance + 1):
        # 提取当前偏移的上对角线
        diag_vals = np.diag(mat, offset=offset)
        # 过滤0元素
        non_zero = diag_vals[diag_vals != 0]
        
        if len(non_zero) == 0:
            np.fill_diagonal(mat[:, offset:], 0)
            continue
        
        # 降序排序并裁剪
        sorted_vals = np.sort(non_zero)[::-1]
        trim_index = max(0, round(trim * len(sorted_vals)) - 1)
        remaining = sorted_vals[trim_index:]
        
        if len(remaining) == 0:
            np.fill_diagonal(mat[:, offset:], 0)
            continue
        
        mu = np.mean(remaining)
        sd = np.std(remaining)
        
        if sd < 1e-6:
            np.fill_diagonal(mat[:, offset:], 0)
        else:
            normalized = (diag_vals - mu) / sd
            # 小于阈值的元素置0
            normalized = np.where(normalized >= mu + sd, normalized, 0)
            np.fill_diagonal(mat[:, offset:], normalized)
    
    # 生成对称矩阵，复制上对角线到下对角线
    mat = mat + mat.T - np.diag(np.diag(mat))
    return mat

额外性能提升建议 #

• 若处理多个矩阵，可将其堆叠为3D数组，沿第三维度做批量向量化操作，避免循环遍历单个矩阵
• 尽量避免在循环内创建新的Python对象，所有操作基于Numpy数组完成
• 若硬件允许，用CuPy替代Numpy，利用GPU加速大矩阵运算，性能可提升数倍

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