基于对角元素由同阶方阵生成新矩阵的NumPy向量化高效实现需求

需求说明 #

给定两个同阶方阵A、B，生成同阶方阵C，其中每个元素C[i,j]的计算规则为：

• 累加一系列项：A[k,k] * B[i-k, j-k]，k从0开始
• 当i-k < 0或j-k < 0时停止累加
• 示例：
  • C[7,4] = A[0,0]*B[7,4] + A[1,1]*B[6,3] + A[2,2]*B[5,2] + A[3,3]*B[4,1] + A[4,4]*B[3,0]（k到4为止，k=5时j-k=-1，停止）
  • C[5,9] = A[0,0]*B[5,9] + A[1,1]*B[4,8] + A[2,2]*B[3,7] + A[3,3]*B[2,6] + A[4,4]*B[1,5] + A[5,5]*B[0,4]（k到5为止，k=6时i-k=-1，停止）

高效向量化实现方案 #

核心思路 #

利用NumPy的stride tricks构造B的所有左上偏移视图（避免数据复制），再与A的对角元素做加权求和，全程无显式循环，最大化利用底层优化。

代码实现 #

import numpy as np

def compute_custom_matrix(A, B):
    n = A.shape[0]
    # 提取A的对角元素
    a_diag = np.diag(A)
    
    # 对B进行补0，构造可生成偏移视图的基础矩阵
    padded_B = np.pad(B, ((0, n), (0, n)), mode='constant')
    # 通过stride tricks生成所有k偏移后的B视图（k从0到n-1）
    # strided[k, i, j] 对应 B[i-k, j-k]，超出原B范围的位置为0
    strided = np.lib.stride_tricks.as_strided(
        padded_B,
        shape=(n, n, n),
        strides=(padded_B.strides[0] + padded_B.strides[1], padded_B.strides[0], padded_B.strides[1])
    )
    
    # 对每个(i,j)，加权求和所有k对应的项
    C = np.tensordot(strided, a_diag, axes=((0,), (0,)))
    # 截取有效区域（原方阵大小）
    return C[:n, :n]

方案优势 #

1. 零数据复制：stride tricks直接创建内存视图，避免额外内存开销
2. 全向量化运算：tensordot调用NumPy底层优化的BLAS/LAPACK接口，远快于Python循环
3. 通用性强：支持任意阶数的方阵，无需针对具体维度修改代码

验证示例 #

假设A是5x5方阵，对角元素为[1,2,3,4,5]，B是5x5随机矩阵：

A = np.diag([1,2,3,4,5])
B = np.random.randint(0, 10, (5,5))
C = compute_custom_matrix(A, B)
# 验证C[3,2] = 1*B[3,2] + 2*B[2,1] + 3*B[1,0]
assert C[3,2] == B[3,2] + 2*B[2,1] + 3*B[1,0]

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