如何用NumPy矩阵运算从邻接矩阵快速生成路径（40000节点）

问题背景 #

给定邻接矩阵（示例如下），需要生成一条遍历节点的路径（示例输出：[0, 2, 1, 3, 4]）。当前使用while循环实现的方法在40000节点的场景下运行缓慢，希望通过矩阵运算的方式提升效率。

示例邻接矩阵：

import numpy as np
a = np.array([[0., 0., 1., 0., 1.],
              [0., 0., 1., 1., 0.],
              [1., 1., 0., 0., 0.],
              [0., 1., 0., 0., 1.],
              [1., 0., 0., 1., 0.]])

现有低效代码：

def create_path_from_joins(joins):
    # not assuming the path is connected
    i = 0
    path = [i]
    elems = np.where(joins[i] == 1)
    elems = elems[0].tolist()
    join_to = set(elems) - set(path)

    while len(join_to) > 0:
        # choose the one that is not already in the path
        elem = list(join_to)[0]
        path.append(elem)
        i = elem
        elems = np.where(np.array(joins[i]) == 1)
        elems = elems[0].tolist()
        join_to = set(elems) - set(path)
    return path

低效原因分析 #

• 循环内频繁调用np.where、列表转换、集合差集运算，这些操作在大节点量下会触发大量内存分配和遍历，时间开销极大
• 集合的差集运算对于40000节点场景，每次操作的时间复杂度都很高
• 逐次遍历的方式没有利用numpy的向量化运算优势，完全依赖Python层面的循环逻辑

矩阵运算优化方案 #

利用numpy的向量化操作替代循环和集合运算，通过布尔数组标记已访问节点，快速定位下一个未访问的邻接节点：

def fast_path_from_adj(adj_matrix):
    n_nodes = adj_matrix.shape[0]
    visited = np.zeros(n_nodes, dtype=bool)
    path = []
    current = 0
    visited[current] = True
    path.append(current)
    
    # 复制邻接矩阵，避免修改原数据
    adj = adj_matrix.copy()
    
    while len(path) < n_nodes:
        # 获取当前节点的所有邻接节点索引
        neighbors = adj[current].nonzero()[0]
        # 筛选出未访问的邻接节点，取第一个
        next_node = neighbors[~visited[neighbors]][0]
        # 更新访问状态和路径
        visited[next_node] = True
        path.append(next_node)
        # 清除双向连接，避免重复遍历已走边（可选，进一步提速）
        adj[current, next_node] = 0
        adj[next_node, current] = 0
        current = next_node
    
    return path

优化点说明 #

• 使用布尔数组visited做访问标记，比集合查找快几个数量级，numpy的布尔索引是底层C实现，效率极高
• 用nonzero()直接提取邻接节点，结合向量筛选替代集合差集运算，全程向量化操作
• 可选的邻接矩阵置0操作，避免每次重复检查已走过的边，减少后续的计算量
• 去掉了循环内的列表转换、集合操作等冗余步骤，充分利用numpy的性能优势

测试示例：调用fast_path_from_adj(a)会返回[0, 2, 1, 3, 4]，与期望结果一致。

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