This code is meant to implement a Graph class which has a method that returns a path between 2 nodes using breadth-first search.

I'm doing this to improve my style and to improve my knowledge of fundamental algorithms/data structures for an upcoming coding interview.

from queue import PriorityQueue
from collections import deque
from copy import deepcopy

class Graph:
    def __init__(self):
        self.nodes = dict()

    def add_node(self, id):
        new_node = Node(id)
        self.nodes[id] = new_node
        return new_node

    def get_node(self, id):
        node = self.nodes[id] if id in self.nodes else None
        return node

    def bfs(self, start, end):
        queue = PriorityQueue()
        queue.put((0, start))
        seen = set([start])
        parents = dict()

        while not queue.empty():
            cur_node = queue.get(0)[1]
            if cur_node is not end:
                for neighbour in cur_node.neighbours:
                    if neighbour not in seen:
                        dist_to_neighbour = cur_node.neighbours[neighbour]
                        queue.put((dist_to_neighbour, neighbour))
                        parents[neighbour] = cur_node
                path = self.get_path(start, end, parents)
                return path
        return None

    def get_path(self, start, end, parents):
        path = deque([end])
        cur_node = end
        while cur_node is not start:
            cur_node = parents[cur_node]
        return list(path)

class Node:
    def __init__(self, id):
        self.id = id
        self.neighbours = dict()

    def add_neighbour(self, to, weight):
        self.neighbours[to] = weight

I have a few ideas about what might be improved but I'm not sure how to do it (or even if it's worth doing).

  1. There might be a better way to do get_node() instead of squashing it into 2 lines like how I've done it.
  2. Maybe cur_node = queue.get(0)[1] should be 2 lines?
  3. Change dist_to_neighbour = cur_node.neighbours[neighbour] to use a method instead of using dictionary notation (which might be a little awkward). Maybe I should have a method instead?
  4. Using something else instead of a deque in get_path() (especially in an interview context where deque might be a bit too exotic).
  5. Have more comments?

Here is my previous attempt at this code. I've also used the answers to my depth-first search attempts here and here to write this code.


1 Answer 1


I don't see that you are using add_node or get_node anywhere, but I'm assuming you would when actually using this class.

Nevertheless, I think these functions would better be called __setitem__ and __getitem__, allowing usage, such as g = Graph(); g[5] = Node(5) and print g[5]. This would also make initializing the graph easier, which leads me to my next recommendation.

I would expect the Graph class to take an argument on initialization to directly build a graph (without having to manually add every single node).

Lastly, get_node can be greatly simplified by using dict.get, which returns None when the key is not present in the dictionary.

class Graph:
    def __init__(self, nodes=None):
        self.nodes = dict()
        if nodes is not None:
            for _id in nodes:
                self[_id] = Node(_id)

    def __setitem__(self, _id, val):
        self.nodes[_id] = val

    def __getitem__(self, _id):
        return self.nodes.get(_id)

    def add_node(self, _id):
        self[_id] = Node(_id)


This allows initializing an instance like this: Graph([1,2,3]). It does not, however, allow initializing a nested structure.

Regarding your questions:

  1. Use dict.get
  2. I would maybe use _, cur_node = queue.get(0) to make it clearer that you don't need the priority here.
  3. I think it is still quite readable.
  4. deque does not seem too exotic to me, but then again, I don't perform job interviews, so who know.

Last but not least, I would have a look at whether you really need queue.PriorityQueue here. One of its features is that it is thread-safe, because it performs locking when adding or popping values from the queue. However, it also incurs a performance hit because of this (I saw a factor 2 being discussed for Python 2.6).

Instead you could roll your own implementation using heapq, similar to the one shown in the implementation details for heapq, maybe without the part about REMOVED entries, because you don't change the priority of a neighbour once you have calculated it. So you could use:

from heapq import heappush, heappop
import itertools

class PriorityQueue:
    def __init__(self, tasks=None):
        '''A new PriorityQueue.
            tasks:    iterable of (priority, task) tuples'''
        self.pq = []                         # list of entries arranged in a heap
        self.counter = itertools.count()     # unique sequence count

        if tasks is not None:
            for priority, task in tasks:
                self.add_task(task, priority)

    def add_task(self, task, priority=0):
        'Add a new task'
        count = next(self.counter)
        entry = (priority, count, task)
        heappush(self.pq, entry)

    def pop_task(self):
        'Remove and return the lowest priority task. Raise KeyError if empty.'
        if self.pq:
            _, _, task = heappop(self.pq)
            return task
        raise KeyError('pop from an empty priority queue')

This should basically be how queue.PriorityQueue is implemented as well, but without the locking. Note that heapq is implemented in C, so it is quite fast. You would have to make some performance comparisons to see how this performs with respect to queue.PriorityQueue.

  • \$\begingroup\$ Thanks for the answer. Using __setitem__ and __getitem__ definitely make it nicer. Being able to initialise a graph with a list makes sense. Your answers to my thoughts were quite helpful. Also, your section about using heapq due to performance was very interesting. I learned a lot of new interesting stuff from that. Thanks :) \$\endgroup\$ Commented Nov 23, 2016 at 6:07

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.