如何向量化内索引受外索引限制的嵌套内积循环？

已知arr1和arr2为形状(N,3)的NumPy数组（N∈[500,2000]），且满足np.dot(arr1[i], arr2[j]) = np.dot(arr1[j], arr2[i])，需要计算所有内积组合，同时尽可能避免冗余计算并提升效率。

核心分析 #

直接调用np.dot(arr1, arr2.T)（或arr1 @ arr2.T）是目前速度最快的方案——这是因为它依赖底层高度优化的BLAS矩阵乘法实现，尽管计算了全部N²个元素，但硬件级并行、缓存友好的内存访问模式带来的性能提升，远超过只计算一半元素的冗余开销。

但如果确实需要严格最小化计算量（仅计算上三角+对角线部分），可以通过向量化的三角索引操作实现，完全避免Python层面的循环。

向量化实现方案 #

利用np.triu_indices一次性获取所有需要计算的上三角（含对角线）索引对，通过广播批量计算内积后填充矩阵，最后对称化得到完整结果：

import numpy as np

N = 2000
arr1 = np.arange(N*3).reshape((N,3))
arr2 = arr1.copy()

# 获取上三角（含对角线）的索引对
i, j = np.triu_indices(N)
# 批量计算内积：利用广播对所有(i,j)对计算点积
inner_vals = np.sum(arr1[i] * arr2[j], axis=1)
# 初始化结果矩阵
inner_products = np.zeros((N, N))
# 填充上三角区域
inner_products[i, j] = inner_vals
# 对称填充下三角，对角线元素需修正（避免重复相加）
inner_products += inner_products.T
inner_products[np.diag_indices(N)] /= 2

方案说明 #

1. 索引批量获取：np.triu_indices(N)直接生成所有i ≤ j的索引对，避免了手动循环生成索引的开销。
2. 向量化计算：通过广播机制，arr1[i]和arr2[j]会自动配对成形状(M,3)的数组（M=N*(N+1)/2），逐元素相乘后求和，一次性完成所有需要的内积计算。
3. 对称填充：通过转置相加快速填充下三角区域，再修正对角线元素（避免被重复计算两次）。

性能对比 #

• 直接矩阵乘法：arr1 @ arr2.T的速度通常是三角索引方案的2_{5倍（N=2000时，矩阵乘法耗时约1}2ms，三角索引方案约5~10ms），原因是BLAS实现的内存访问连续性更强，缓存命中率更高，且充分利用CPU向量指令集。
• 三角索引方案：虽然计算量仅为矩阵乘法的一半，但内存访问模式更分散，缓存效率较低，因此整体速度不如前者，但比手动嵌套循环（包括部分向量化的单循环版本）快一个数量级以上。

结论 #

• 若追求极致速度，优先选择arr1 @ arr2.T——尽管计算了全部元素，但底层优化的收益远大于冗余计算的成本。
• 若必须最小化计算量，使用np.triu_indices的向量化方案，完全规避Python循环，效率远高于手动嵌套实现。

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