I previously attempted to implement a Graph class which has a method that returns a path between 2 nodes using depth-first search. This is my latest attempt at this (using the feedback I got from last time).

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

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

    def add_node(self, id):
        new_node = Node(id)
        return new_node

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

    def dfs(self, start, end):
        stack = [start]
        parents = {}
        seen = set([start])

        while stack:
            cur_node = stack.pop()
            if cur_node is not end:
                for neighbour in cur_node.neighbours:
                    if neighbour not in seen:
                        parents[neighbour] = cur_node
                path = self.get_path(start, cur_node, parents)
                return path
        return None

    def get_path(self, start, end, parents):
        stack = [end]
        cur_node = end
        while cur_node is not start:
            cur_node = parents[cur_node]

        path = []
        while stack:
            cur_node = stack.pop()
        return path

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

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

Honestly your code is pretty good.

There are however two things that I would change:

  1. self.nodes to be a dictionary rather than a set. This allows you to remove the need for get_node, and changes it to \$O(1)\$ from \$O(n)\$.

  2. In get_path you pop from one stack and append to another. Instead of doing this manually I'd use the built-in list slicing notation, stack[::-1].

  • \$\begingroup\$ Thanks for the answer. I like the suggestion of using a dictionary instead of a set. That simplifies it even more. Your suggestion about just iterating backwards is also good. What do you think about using for n in reversed(stack): instead of stack[::-1]? \$\endgroup\$ Nov 22 '16 at 7:16
  • \$\begingroup\$ What benefits do you get for not using stack[::-1]? I can't think of any, but it costs readability. \$\endgroup\$
    – Peilonrayz
    Nov 22 '16 at 8:53
  • \$\begingroup\$ I guess it's just a matter of opinion. I personally find for n in reversed(stack) to be a bit more readable. But stack[::-1] is also good :) \$\endgroup\$ Nov 22 '16 at 8:55

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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