set
s perform containing checks (w in visited
) \$O(1)\$ rather than \$O(n)\$ for lists.
collections.deque
are better than lists for poping elements at the front (popleft
).
- you should put your example code under an
if __name__ == '__main__'
clause.
w
as a variable name does not convey meaning, you should try to come up with something more explicit.
import collections
def breadth_first_search(graph, root):
visited, queue = set(), collections.deque([root])
while queue:
vertex = queue.popleft()
for neighbour in graph[vertex]:
if neighbour not in visited:
visited.add(neighbour)
queue.append(neighbour)
if __name__ == '__main__':
graph = {0: [1, 2], 1: [2], 2: []}
breadth_first_search(graph, 0)
Given a growing number of comments indicating that the code does not return anything, I’d like to add that, yes, this code does not process nodes: it only traverse the graph and you're likely to want to add your own custom logic to process each node. As your mileage may vary (building a traversal list, finding the first node that satisfies a condition, etc.), there is not a "one code fits all" approach, but a useful first approximation would be to yield
each node as they are traversed:
import collections
def breadth_first_search(graph, root):
visited, queue = set(), collections.deque([root])
while queue:
vertex = queue.popleft()
yield vertex
visited.add(vertex)
queue.extend(n for n in graph[vertex] if n not in visited)
if __name__ == '__main__':
graph = {1: [2, 4, 5], 2: [3, 6, 7], 3: [], 4: [], 5: [], 6: [], 7: []}
list(breadth_first_search(graph, 1)) # [1, 2, 4, 5, 3, 6, 7]
Note that this alternative iteration also takes care of the bug mentioned in this answer