Interview question for recurse center. Depth First Search in Python.

rc_interview_dfs.py

import string

def dfs(tree, element):
    """An iterative implementation of depth first search.  It returns the element in question
    if it is present and none otherwise.
    """
    visited = set()
    stack = [tree]
    
    while stack:
        current_node = stack.pop()
        visited.add(current_node)
        if current_node.data == element:
            return element
        elif current_node.children is not None:
            for node in current_node.children:
                if node not in visited:
                    stack.append(node)
        
def bfs(element):
    """Breadth First Search TODO::"""
    pass


class Node(object):
    
    def __init__(self, children = None, data=None):
        self.children = children
        self.data = data
        
    def add_child(self, node):
        if not self.children:
            self.children = [node]
        else:
            self.children.append(node)
            
    def add_children(self, nodes):
        if not self.children:
            self.children = nodes
        else:
            self.children += nodes
            
def create_test_tree():
    """This method creates a tree with all of the letters in the alphabet present for testing purposes"""
    root = Node(data='a')
    current_node = root
    for i, letter in enumerate(string.ascii_lowercase[1:]):
        new_tree = Node(data=letter)
        current_node.add_child(new_tree)
        if i == 13:
            current_node = root
        elif i % 2 == 0:
            current_node = new_tree
        
    return root
  
tree = create_test_tree()
print(dfs(tree, 'a'))
print(dfs(tree, 'b'))
print(dfs(tree, 'w'))
print(dfs(tree, 'aaa'))