You could keep every item as (node, distance)
and append (neighbour, distance+1)
import matplotlib.pyplot as plt
import networkx as nx
# --- functions ---
def bfs(g, node):
distance = 0
visited = [node]
queue = [(node, distance)]
tree = {}
while queue:
s, distance = queue.pop(0)
tree[s] = []
for neighbour in list(g.adj[s]):
if neighbour not in visited:
visited.append(neighbour)
tree[s].append((neighbour, distance+1))
queue.append((neighbour, distance+1))
return tree
# --- main ---
#G = nx.erdos_renyi_graph(12, .1)
G = nx.Graph()
G.add_edges_from(([1,2],[2,3], [2,4], [3, 5], [4,5]))
tree = bfs(G, 1)
print(tree)
nx.draw(G, with_labels=True)
plt.show()
Result (manually reformated):
{
1: [(2, 1)],
2: [(3, 2), (4, 2)],
3: [(5, 3)],
4: [],
5: []
}
For graph
Evenutally you can even use tree[(s, distance)]
def bfs(g, node):
distance = 0
visited = [node]
queue = [(node, distance)]
tree = {}
while queue:
s, distance = queue.pop(0)
tree[(s, distance)] = []
for neighbour in list(g.adj[s]):
if neighbour not in visited:
visited.append(neighbour)
tree[(s, distance)].append((neighbour, distance+1))
queue.append((neighbour, distance+1))
return tree
Result (manually reformated):
{
(1, 0): [(2, 1)],
(2, 1): [(3, 2), (4, 2)],
(3, 2): [(5, 3)],
(4, 2): [],
(5, 3): []
}
BTW:
I was thinking to use json.dumps(tree, indent=2)
to reformat result automatically but it can't convert node
to json
.
But pretty printer
can format it in similar way
import pprint
pprint.pprint(tree, indent=2)
Result:
{ (1, 0): [(2, 1)],
(2, 1): [(3, 2), (4, 2)],
(3, 2): [(5, 3)],
(4, 2): [],
(5, 3): []}