如何获取指定最大层级的图节点邻域集合？

我需要获取所选节点邻域内的所有节点。例如：

• 指定层级为1时，结果应为 [black, red, red, red, red]
• 指定层级为2时，结果应为 [black, red, red, red, red, green, green, green, green, green, green, green, green]

我尝试用递归实现，但它的运行方式类似DFS，且由于停止条件设置为depth <= max_level，导致部分应包含的节点被遗漏。

以下是我编写的代码：

def get_neighborhood(node, edges, neighborhood_level, depth, nodes_acc):
    neighbors_filter = np.where(edges[1] == node)
    all_neighbors    = edges[0][neighbors_filter]
    neighbors        = all_neighbors[check_if_not_in(all_neighbors, nodes_acc)]

    if depth == neighborhood_level:
        return nodes_acc

    for neighbor in neighbors:
        nodes_acc = torch.unique(torch.cat((nodes_acc, torch.Tensor([neighbor]))))
    
    for neighbor in neighbors:
        nodes_acc = torch.unique(torch.cat((nodes_acc, get_neighborhood(neighbor, edges, neighborhood_level, depth + 1, nodes_acc))))

    return nodes_acc

参数说明 #

• node：

  tensor(some_id)
  

• edges：

  tensor([[source, source, source, ..., source, source],
          [target, target, target, ..., target, target]])
  

希望得到相关改进建议，输入数据不局限于torch.Tensor，numpy.array或普通列表均可，我后续会自行转换格式。谢谢！

内容的提问来源于stack exchange，提问作者jan-swiatek