Python循环获取各node_id邻居：如何整理为嵌套列表输出

看起来你现在的代码是把所有节点的邻居都收集到了一个扁平列表里，想要改成按每个node_id分组的嵌套列表，其实有两种简单的方式可以实现：

方法一：直接修改原循环（推荐） #

既然你本来就有每个节点对应的邻接表项，不如直接在遍历的时候就构建嵌套列表，这样不用先生成扁平列表再处理，效率更高：

用for循环实现 #

new_list = []
# 遍历每个node_id对应的索引
for i in range(len(node_id)):
    # 为当前节点初始化一个空的邻居列表
    current_neighbours = []
    for arc in self._adjacency_lists[i]:
        current_neighbours.append(arc.head_node_id)
    # 把当前节点的邻居列表加入最终结果
    new_list.append(current_neighbours)

用while循环实现（对应你原来的循环结构） #

如果你更习惯用while循环，也可以改成这样：

new_list = []
i = 0
while i < len(node_id):
    current_neighbours = []
    for arc in self._adjacency_lists[i]:
        current_neighbours.append(arc.head_node_id)
    new_list.append(current_neighbours)
    i += 1

这样运行后，new_list就是你想要的[[1,2], [2], [3], [1], [3]]（对应你描述的节点邻居分布）。

方法二：从已有的扁平列表转换 #

如果你已经有了那个扁平的neighbours列表，并且知道每个节点对应的邻居数量（比如node0有2个，node1有1个，node2有1个，node3有1个，node4有1个），可以通过切片来分割：

# 假设每个节点的邻居数量依次是[2,1,1,1,1]
counts = [2,1,1,1,1]
new_list = []
start = 0
for cnt in counts:
    new_list.append(neighbours[start:start+cnt])
    start += cnt

不过这种方法需要你提前知道每个节点的邻居数量，不如第一种方法直接从源数据构建来得稳妥。

更Pythonic的简化写法 #

还可以用嵌套列表推导式把代码压缩成一行，逻辑和上面的for循环完全一致：

new_list = [[arc.head_node_id for arc in self._adjacency_lists[i]] for i in range(len(node_id))]

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