如何向量化该for循环？numpy分块矩阵乘法实现

好问题！这种分块运算的向量化实现能帮你避开低效的Python循环，充分利用numpy的底层优化。下面我会针对等长度划分和不等长度划分两种场景给出具体方案：

场景1：等长度划分（最常见且高效） #

假设你的数组f长度n正好是划分份数r的整数倍，即每个子块fᵢ的长度为k = n // r，矩阵A的每块Aᵢ对应k×m的子矩阵。这种情况可以直接通过reshape结合矩阵乘法一步到位：

import numpy as np

# 示例数据
n = 6  # f的长度
m = 4  # A的列数
r = 3  # 划分份数
k = n // r  # 每个子块的长度

f = np.random.rand(n)          # shape: (6,)
A = np.random.rand(n, m)       # shape: (6, 4)

# 向量化实现
f_reshaped = f.reshape(r, k)   # 将f拆分为r个长度为k的行向量，shape: (3, 2)
A_reshaped = A.reshape(r, k, m)# 将A拆分为r个k×m的子矩阵，shape: (3, 2, 4)
B = f_reshaped @ A_reshaped    # 分块乘法，自动广播，shape: (3, 4)

验证正确性 #

用循环实现对比结果，确认向量化输出和手动分块计算一致：

# 循环计算作为基准
B_loop = np.zeros((r, m))
for i in range(r):
    start = i * k
    end = (i + 1) * k
    B_loop[i] = f[start:end] @ A[start:end]

print(np.allclose(B, B_loop))  # 输出True，说明结果一致

场景2：不等长度划分 #

如果各子块长度不相等（比如n=7，r=3，子块长度为[2,3,2]），我们可以通过构造掩码矩阵实现向量化，避免显式循环：

# 示例数据
n = 7
m = 4
r = 3
k_list = [2, 3, 2]  # 每个子块的长度列表，总和为n

f = np.random.rand(n)          # shape: (7,)
A = np.random.rand(n, m)       # shape: (7, 4)

# 向量化构造掩码矩阵
block_ids = np.repeat(np.arange(r), k_list)  # 标记每个元素所属的块，shape: (7,)
mask = np.zeros((r, n), dtype=f.dtype)
np.add.at(mask, (block_ids, np.arange(n)), f)  # 向量化赋值，将f的元素放到对应块的行

# 计算最终结果
B = mask @ A  # shape: (3, 4)

验证正确性 #

同样用循环对比：

B_loop = np.zeros((r, m))
start = 0
for i in range(r):
    end = start + k_list[i]
    B_loop[i] = f[start:end] @ A[start:end]
    start = end

print(np.allclose(B, B_loop))  # 输出True，结果一致

为什么这两种方法是向量化的？ #

• 等长度划分中，reshape和@都是numpy的底层优化操作，完全避开了Python级别的循环，直接调用C实现的运算逻辑。
• 不等长度划分中，np.repeat和np.add.at都是向量化的数组操作，掩码矩阵与A的乘法也是高度优化的矩阵运算，效率远高于手动循环。

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