基于Numpy优化Python嵌套For循环：二次分配问题局部搜索算法调优

我之前在优化二次分配问题（QAP）的局部搜索算法时，也被三层嵌套Python循环的性能问题折磨过——CPython解释器处理这类循环的效率实在拉胯，尤其是数据规模上去之后，运行时间直接爆炸。不过好在Numpy的向量化和广播特性天生就是用来解决这类问题的，下面分享几个我亲测有效的优化思路：

1. 用Numpy广播彻底压平嵌套循环 #

三层循环的核心问题是Python层面的逐元素遍历，而Numpy的广播机制可以让我们把这些循环转换成底层C实现的矩阵运算，速度能提升几个数量级。

举个例子，假设你原来的循环是遍历solution中的节点对（i,j），再遍历所有k来计算交换i和j后的代价变化，那完全可以把solution转换成合适的形状，让Numpy自动完成所有元素的并行计算。比如把solution变成列向量后，和自身的行向量组合，就能一次性生成所有节点对的距离/流量关联矩阵，再和预定义的flow/dist矩阵做点积求和，直接替代k层的循环。

2. 用爱因斯坦求和（np.einsum）简化复杂的多维度计算 #

如果你的循环里涉及到多个矩阵的乘积和求和，np.einsum绝对是神器。它可以用简洁的下标符号描述复杂的求和逻辑，比嵌套np.dot或np.sum更直观，而且性能也不差。

比如原来三层循环里的累加计算：

total = 0
for i in range(n):
    for j in range(n):
        for k in range(n):
            total += flow[i,k] * dist[solution[j], solution[k]]

可以直接写成：

total = np.einsum('ik,jk->', flow, dist[solution][:, solution])

一行代码替代三层循环，效率直接拉满。

3. 预计算重复使用的矩阵，避免循环内冗余计算 #

局部搜索算法中很多计算是重复的，比如当前solution对应的距离矩阵dist[solution][:, solution]，完全可以在循环外提前计算好，而不是每次遍历都重新索引。这样能减少大量不必要的内存访问和计算开销。

另外，如果你的mask是用来过滤某些节点的，也可以提前生成掩码矩阵，用布尔索引直接过滤掉无效的节点对，不用在循环里一次次判断mask[i]或mask[j]。

4. 批量计算所有可能的操作，替代逐次迭代评估 #

局部搜索通常需要评估多个候选解（比如所有可能的2-opt交换对），与其逐个遍历计算每个交换的代价，不如一次性算出所有候选解的代价变化，再从中选最优的。比如用广播生成所有节点对的代价变化矩阵，然后直接取最小值对应的交换对，这样整个过程都是向量/矩阵运算，完全没有Python循环。

简单示例：从循环到向量化的转换 #

假设你原来的三层循环代码大概是这样的（计算交换节点i和j的代价变化）：

n = len(solution)
delta = np.zeros((n, n))
for i in range(n):
    if not mask[i]:
        continue
    for j in range(n):
        if i == j or not mask[j]:
            continue
        cost_change = 0
        for k in range(n):
            if k == i or k == j:
                continue
            cost_change += flow[i,k] * (dist[solution[j], solution[k]] - dist[solution[i], solution[k]])
            cost_change += flow[k,i] * (dist[solution[k], solution[j]] - dist[solution[k], solution[i]])
            cost_change += flow[j,k] * (dist[solution[i], solution[k]] - dist[solution[j], solution[k]])
            cost_change += flow[k,j] * (dist[solution[k], solution[i]] - dist[solution[k], solution[j]])
        delta[i,j] = cost_change

优化后的向量化版本：

n = len(solution)
# 预计算当前solution对应的距离矩阵
current_dist = dist[solution][:, solution]

# 生成有效节点对的掩码（排除i==j和mask不允许的节点）
valid_mask = np.outer(mask, mask) & (np.eye(n) == 0)

# 计算交换i,j后的距离变化矩阵
dist_diff_i = dist[solution[:, None], solution] - current_dist[:, None, :]
dist_diff_j = dist[solution, solution[:, None]] - current_dist[None, :, :]

# 用einsum计算所有i,j,k的求和
delta_i = np.einsum('ik,ijk->ij', flow, dist_diff_i) + np.einsum('ki,kij->ij', flow.T, dist_diff_i.transpose(1,0,2))
delta_j = np.einsum('jk,jik->ij', flow, dist_diff_j.transpose(1,0,2)) + np.einsum('kj,kji->ij', flow.T, dist_diff_j)

delta = delta_i + delta_j
# 过滤掉无效的节点对
delta[~valid_mask] = np.inf  # 或者设为一个极大值，避免后续选中

这样一来，所有的循环都被转换成了Numpy的底层运算，速度会有质的飞跃。

最后提醒一句：优化过程中要多打印数组的形状（print(arr.shape)），确保广播的维度是正确的，避免出现形状不匹配的错误。如果遇到复杂的维度转换，np.expand_dims可以帮你灵活调整数组的形状。

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