Python+Networkxでグラフ上のある頂点からKステップ分の近傍を簡単に探す

以前の凸包や最短経路を求める例と似た感じの小ネタで，今回はある頂点からK-th order neighborを探してくるスクリプトを，できるだけ公式の機能だけを使ってやってみる．もっと簡単な方法があったらコソっと教えてください．とりあえずの例題として，幾何的なグラフ構造をランダムに作成した．

ランダムなグラフの作成

import networkx as nx
import matplotlib.pyplot as plt

G = nx.random_geometric_graph(200, 0.125)
pos = nx.get_node_attributes(G, 'pos')

# find node near center (0.5,0.5)
dmin = 1
ncenter = 0
for n in pos:
    x, y = pos[n]
    d = (x - 0.5) ** 2 + (y - 0.5) ** 2
    if d < dmin:
        ncenter = n
        dmin = d
p = nx.single_source_shortest_path_length(G, ncenter)

fig = plt.figure(figsize=(8,8))
nx.draw_networkx_edges(G, pos, nodelist=[ncenter],alpha=0.4)
nx.draw_networkx_nodes(G, pos, node_size=40)
plt.xlim(-0.05,1.05)
plt.ylim(-0.05,1.05)
plt.axis('off')
plt.show()

ある頂点からKステップの頂点を求める

今回は最短経路を求める関数に渡すことができる引数cutoffを利用した．

K = 5
nodedict = {}
for k in range(1, K + 1):
    Nk = nx.single_source_shortest_path_length(G, ncenter, cutoff=k)
    Nk = [n for (n, v) in Nk.items() if v == k]
    nodedict[k] = Nk

可視化してみる

fig = plt.figure(figsize=(8,8))
nx.draw_networkx_edges(G, pos, nodelist=[ncenter], alpha=0.4)
nx.draw_networkx_nodes(G, pos, node_size=40, node_color='black')
nx.draw_networkx_nodes(G, pos, node_size=50, node_color='red', nodelist=nodedict[1])
nx.draw_networkx_nodes(G, pos, node_size=50, node_color='green', nodelist=nodedict[2])
nx.draw_networkx_nodes(G, pos, node_size=50, node_color='cyan', nodelist=nodedict[3])
nx.draw_networkx_nodes(G, pos, node_size=50, node_color='blue', nodelist=nodedict[4])
nx.draw_networkx_nodes(G, pos, node_size=50, node_color='magenta', nodelist=nodedict[5])
plt.xlim(-0.05,1.05)
plt.ylim(-0.05,1.05)
plt.axis('off')
plt.show()