如何向量化对称矩阵归一化操作以提升numpy运算效率？

原代码通过双重Python循环遍历矩阵上三角元素完成归一化，当N≈3800时耗时约12秒，核心问题是Python循环的高开销——对于O(N²)量级的元素操作，纯Python循环的效率远低于NumPy底层的C实现。以下是两种高效优化方案：

方案一：全向量化操作（推荐） #

利用NumPy的广播机制，将整个矩阵的归一化操作转化为一次向量外积运算，完全规避Python循环：

import numpy as np

def gen_new_mat_fast(mat: np.ndarray):
    # 提取对角线元素，计算每个元素的倒数平方根
    diag_reciprocal_sqrt = 1.0 / np.sqrt(np.diag(mat))
    # 通过广播实现逐元素归一化：A[x,y] = A[x,y] / sqrt(A[x,x]*A[y,y])
    # 等价于 A[x,y] * (1/sqrt(A[x,x])) * (1/sqrt(A[y,y]))
    new_mat = mat * diag_reciprocal_sqrt[:, np.newaxis] * diag_reciprocal_sqrt[np.newaxis, :]
    return new_mat

优势 #

• 底层由C实现的NumPy内核执行，无Python循环开销，速度可提升数十倍甚至上百倍；
• 利用矩阵对称性，无需单独处理下三角元素，结果自然保持对称；
• 代码简洁，可读性强。

方案二：原地修改（内存优化版） #

若矩阵规模极大、内存紧张，可选择原地修改原矩阵（注意：会改变输入矩阵），进一步节省内存开销：

import numpy as np

def gen_new_mat_inplace(mat: np.ndarray):
    diag_reciprocal_sqrt = 1.0 / np.sqrt(np.diag(mat))
    # 分两步原地完成归一化
    mat *= diag_reciprocal_sqrt[:, np.newaxis]
    mat *= diag_reciprocal_sqrt[np.newaxis, :]
    return mat

注意事项 #

• 若需保留原矩阵，需先执行new_mat = mat.copy()再进行原地操作；
• 确保输入矩阵的对角线元素均为正数，否则np.sqrt会返回NaN或复数（原代码已隐含此前提）。

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