如何向量化满足j < i+1条件的双重求和运算？

针对你提到的「仅对j < i+1（也就是j ≤ i，假设i、j为整数索引）的(i,j)组合进行双重求和」的需求，我们可以利用NumPy的向量化工具替代双重循环，大幅提升计算效率。下面分步骤讲解具体实现思路和示例：

核心思路：定位满足条件的索引对 #

j < i+1等价于j的取值范围是从起始索引到i，对应二维索引中的「下三角区域（含对角线）」。我们可以通过NumPy提供的工具直接定位这些索引，或者生成掩码筛选目标元素，再进行求和。

具体实现示例 #

假设我们的求和式是$\sum_{i=0}^{{n-1}\sum_{j=0}}{i} f(i,j)$，下面用几种常见场景演示：

场景1：f(i,j)是i和j的简单函数（如i+j） #

原双重循环代码：

import numpy as np

n = 5
total = 0
for i in range(n):
    for j in range(i+1):  # j < i+1 → j从0到i
        total += i + j
print(total)  # 输出：50

向量化实现方式1：用tril_indices直接获取目标索引

np.tril_indices(n)会返回两个数组，分别对应下三角区域所有元素的i和j索引，完美匹配j ≤ i的条件：

n = 5
i_indices, j_indices = np.tril_indices(n)
total = np.sum(i_indices + j_indices)
print(total)  # 输出：50

向量化实现方式2：生成网格+掩码筛选

先通过meshgrid生成i和j的二维网格，再用掩码筛选出满足j ≤ i的元素：

n = 5
i_grid, j_grid = np.meshgrid(np.arange(n), np.arange(n), indexing='ij')
# 创建掩码：只保留j ≤ i的位置
mask = j_grid <= i_grid
total = np.sum((i_grid + j_grid)[mask])
print(total)  # 输出：50

向量化实现方式3：利用np.tril生成下三角矩阵求和

np.tril(A)会将数组A的上三角部分置为0，直接对整个矩阵求和即可得到目标结果：

n = 5
i_grid, j_grid = np.meshgrid(np.arange(n), np.arange(n), indexing='ij')
A = i_grid + j_grid
total = np.sum(np.tril(A))
print(total)  # 输出：50

场景2：f(i,j)基于两个一维数组（如x[i]*y[j]） #

假设我们有两个一维数组x和y，需要计算$\sum_{i=0}^{{n-1}\sum_{j=0}}{i} x[i] \times y[j]$：

原双重循环代码：

x = np.array([1,2,3,4,5])
y = np.array([10,20,30,40,50])
total = 0
for i in range(len(x)):
    for j in range(i+1):
        total += x[i] * y[j]
print(total)  # 输出：1900

向量化实现方式1：索引对直接计算

x = np.array([1,2,3,4,5])
y = np.array([10,20,30,40,50])
i_indices, j_indices = np.tril_indices(len(x))
total = np.sum(x[i_indices] * y[j_indices])
print(total)  # 输出：1900

向量化实现方式2：矩阵乘法优化（大数组更高效）

可以构造一个下三角全1矩阵，通过矩阵乘法实现批量计算，这种方式利用了NumPy优化过的线性代数运算，适合大尺寸数组：

x = np.array([1,2,3,4,5])
y = np.array([10,20,30,40,50])
# 生成下三角全1矩阵
tril_mat = np.tril(np.ones((len(x), len(y))))
# x点乘下三角矩阵，再点乘y，得到求和结果
total = x @ tril_mat @ y
print(total)  # 输出：1900

为什么向量化更高效？ #

Python的循环是解释执行的，每一次循环都有额外的开销；而NumPy的向量化操作是在底层用C实现的，能充分利用CPU的并行计算能力，当n很大（比如上万级）时，速度提升会非常显著。

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