寻求纯NumPy实现Watts-Strogatz算法的无循环优化方案

你已经完成了Watts-Strogatz（WS）随机图生成的大部分NumPy实现，但重连阶段的循环操作想替换为纯NumPy向量化逻辑，以下是可行的实现思路及优化代码：

核心问题分析 #

原循环的核心操作是：对每条待删除的边(i,j)，从i的候选池选新边(i,j_new)，然后更新邻接矩阵和候选池（移除新边、加回原边）。难点在于候选池是动态变化的，且每个节点的候选池长度不一致，无法直接用单一的向量化操作覆盖所有情况，但可以通过「批量预处理选择结果 + 批量更新」的方式替代循环。

优化思路 #

1. 批量获取待重连边索引 #

先一次性提取所有需要重连的边的起点i和原终点j，避免循环中反复查询：

i_indices, j_indices = np.nonzero(edges_to_be_removed)
M = len(i_indices)

2. 预生成所有新边的目标节点 #

由于每个节点的候选池（A_inv_mask中i行的非零元素）长度不同，无法用单一的np.random.choice直接生成所有j_new，但可以用列表推导批量处理每个i的候选池，再随机选择：

candidate_lists = [np.where(A_inv_mask[i])[0] for i in i_indices]
j_new_indices = np.array([np.random.choice(cands) for cands in candidate_lists])

这一步的列表推导相比原循环，只是把单条边的选择逻辑批量执行，效率更高。

3. 批量更新邻接矩阵与候选池 #

得到所有j_new后，用NumPy的索引广播特性批量完成所有更新操作，替代循环中的单条边修改：

• 批量添加新边到邻接矩阵
• 批量从候选池移除新边的双向条目
• 批量将原删除边的双向条目加回候选池

优化后的完整代码 #

import numpy as np

def WS_gen(A, N, k, p):
    """
    N : 节点数量
    k : 一维环形图的配位数
    p : 重连概率
    A : 符合N、k参数的有效环形图邻接矩阵
    """

    # 仅考虑上三角矩阵，确保每条边只处理一次
    A_u = np.triu(A)

    # 找出待删除边的矩阵并执行删除
    x = np.random.rand(*A.shape)
    x = np.where(x > (1-p), 1, 0)
    edges_to_be_removed = np.array(x*A_u, dtype=int)
    A_u = A_u - edges_to_be_removed
    
    # 初始化候选池掩码（过滤原边和自环）
    A_inv = np.where(A > 0, 0, 1)
    A_inv = A_inv - np.diag(np.diag(A_inv))
    A_inv_mask = A_inv.copy()
    
    # 获取所有待重连边的索引
    i_indices, j_indices = np.nonzero(edges_to_be_removed)
    M = len(i_indices)
    
    # 无待重连边时直接返回
    if M == 0:
        return A_u + A_u.T
    
    # 批量为每个待重连节点选择新的目标节点
    candidate_lists = [np.where(A_inv_mask[i])[0] for i in i_indices]
    j_new_indices = np.array([np.random.choice(cands) for cands in candidate_lists])
    
    # 批量更新邻接矩阵与候选池
    # 添加新边到上三角邻接矩阵
    A_u[i_indices, j_new_indices] = 1
    # 移除候选池中的新双向边
    A_inv_mask[i_indices, j_new_indices] = 0
    A_inv_mask[j_new_indices, i_indices] = 0
    # 将原删除的双向边加回候选池
    A_inv_mask[i_indices, j_indices] = 1
    A_inv_mask[j_indices, i_indices] = 1
    
    # 恢复对称邻接矩阵
    return A_u + A_u.T

额外说明 #

• 若想彻底避免列表推导，可以考虑预先生成每个节点的候选池索引数组，再用np.random.choice的size参数批量选择，但由于候选池长度不一致，这种方式需要额外处理索引对齐问题，反而会增加代码复杂度，不如列表推导简洁高效。
• 该优化版本保留了原算法的逻辑正确性，同时大幅提升了重连阶段的执行效率，尤其在节点数量较多时效果明显。

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