如何用Numpy高效计算邻接矩阵的路径长度？（TSP遗传算法优化）

你正在优化旅行商问题（TSP）的遗传算法，其中路径长度计算环节耗时最长，当前使用的代码如下：

from itertools import pairwise
import numpy as np
from random import shuffle

def get_path_len(adj_mat: np.ndarray, path: np.ndarray) -> float:
  return sum(adj_mat[i, j] for i, j in pairwise(path)) + adj_mat[path[-1], path[0]]

mat = np.random.randint(1, 1000, (100, 100))
path = np.asarray(list(range(100)))
shuffle(path)

print(get_path_len(mat, path))

最高效的Numpy实现方法 #

原方法依赖Python层面的pairwise迭代和生成器求和，存在大量循环开销，而Numpy的核心优势是矢量化操作，可以完全规避Python循环，大幅提升计算效率。

高效实现的思路是直接构造闭环路径的节点对索引，利用Numpy的数组索引一次性提取所有需要的邻接矩阵元素，再求和：

import numpy as np
from random import shuffle

def get_path_len_fast(adj_mat: np.ndarray, path: np.ndarray) -> float:
    # 构造目标节点索引：path的下一个节点，最后一个节点指向第一个节点
    idx_j = np.concatenate([path[1:], path[:1]])
    # 利用Numpy的二维数组索引，一次性获取所有路径段的长度并求和
    return adj_mat[path, idx_j].sum()

# 测试代码
mat = np.random.randint(1, 1000, (100, 100))
path = np.asarray(list(range(100)))
shuffle(path)

print(get_path_len_fast(mat, path))

效率提升原因 #

• 原方法中pairwise和生成器求和是Python级别的循环，每个元素的访问和累加都有解释器开销；
• 优化后的方法完全基于Numpy的底层C实现操作：数组索引、拼接、求和都是矢量化执行，没有Python循环的额外消耗，当路径节点数越多（比如TSP问题中常见的数百上千节点），性能差距会越显著。

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