DFS function with class Tree in Python

tree_and_dfs.py

# Recurse Center Application 2016, Emily Boynton
# Code below includes a class Tree which can create a navigable tree
# and a function dfs() or "Depth-first Search" which allows a user to 
# search the entire tree for a specific tree (by its name) 


# class Tree gives each "node" a name, array of children, and parent. 
# Tree navigation includes: parent(), find_child(), to_root(), add()
# Tree functions include: print_children(), print_whole_tree(), print_tree()
class Tree:
    def __init__(self, name=None):
        self.name = name
        self.children = []
        self.parent = None


    # Returns the parent tree of the current tree
    def parent(self):
        return self.parent


    # Searches and navigates to child tree of current tree
    def find_child(self, search_name):
        for child in self.children:
            if child.name == search_name:
                return child


    # Navigates to the root tree from the current tree by returning 
    # parent until it reaches parent-less tree
    def to_root(self):
        if self.parent == None:
            return self
        return self.parent.to_root()


    # Adds tree to current tree by appending a node to the current children array.
    # Also assigns self as parent tree to new child
    def add(self, name, children=None):
        new_tree = Tree(name)
        self.children.append(new_tree)
        new_tree.parent = self  
        return new_tree


    # Prints all children of the current tree
    def print_children(self):
        for child in self.children:
            print child.name


    # Print whole tree ensures that root tree is passed to print_tree()
    def print_whole_tree(self):
        # If tree has a parent, navigate to root tree
        if self.parent != None:
            node_to_print = self.to_root()
        # Tree is root tree
        else:
            node_to_print = self
        self.print_tree(node_to_print)


    # Print tree is called by print_whole_tree()
    # Recursively prints tree with formatting
    def print_tree(self, node_to_print):
        print node_to_print.name, ": ",

        for child in node_to_print.children:
            print child.name,
        print ""
        for child in node_to_print.children:
            child.print_tree(child)
        return


# Depth-first Search function which accepts a tree and a node name 
# Return the node if it exists, and None if not.
def dfs(tree, search_name):
    # If the name matches the search query, return the tree
    if tree.name == search_name:
        return tree

    # If the tree has no children, return false
    if tree.children == []:
        return False

    # Tree has children    
    else:
        # For each child in the children, recursively call dfs() to find if match
        for child in tree.children:
            match = dfs(child, search_name)
            # If a match is found, return the node by calling dfs()
            if match:
                return dfs(child, search_name)

    # No match found: print message and return a NoneType
    print "no match!"
    return None

tree_and_dfs_update.py

# Recurse Center Application 2016, Emily Boynton
# Code below includes a class Tree which can create a navigable tree
# and a function dfs() or "Depth-first Search" which allows a user to 
# search the entire tree for a class tree 


# class Tree gives each "node" a name, array of children, and parent. 
# Tree navigation includes: parent(), find_child(), to_root(), add_child()
# Tree functions include: print_children(), print_whole_tree(), print_tree()


# UPDATES FOR MAX 5/3
# Main updates include: 
#   - simplified/fixed-up depth-first search!
#   - added a tree_dict to keep a list of trees
#   - added a self.visited flag to the class Tree
#   - added set_tree_false() which preps the visited flag for a search
#   - there is now a dfs() which calls set_tree_false() and the recursive depth(). 
#   - no longer use vocabulary of parent/child only have a list of self.points_to
#   - got rid of unnecessary Tree functions
#   - added some cycles to sample RC tree ("h" --> "a", "c" --> "c")
#   - added a main function so that you can more readily see my testing


# tree_dict will hold an unordered list of the name of a tree and its instance
tree_dict = {}

class Tree:
    def __init__(self, name=None):
        self.name = name
        self.points_to = []
        # adds the Tree name and instance as a key-value pair to tree_dict
        tree_dict[name] = self
        # default value of false. Utilized as flag in dfs()
        self.visited = False


    # Searches the current tree and moves to the searched tree
    # if it directly points to it
    # if no match, return itself
    def move_to_tree(self, search_name):
        for tree in self.points_to:
            if tree.name == search_name:
                return tree
        return self 


    # Appends another tree to self.points_to
    # First determining if the tree is already in the graph structure
    # If yes, point to that tree
    # Otherwise, add a new tree and then add it to self.points_to
    def add_tree(self, name, children=None):
        # call search on name, to see if in structure
        if search(name):
            # if it is, append it to self.points_to
            self.points_to.append(search(name))
            # return the found tree
            return tree_dict[name]
        # else create and append a new_tree to self.points_to
        else:
            new_tree = Tree(name)
            self.points_to.append(new_tree) 
            # return the new_tree
            return new_tree


    # Prints all of the trees pointed to by the current tree
    def print_directed_trees(self):
        print self.name, ": ",
        # print the trees pointed to by each tree
        for tree in self.points_to:
            print tree.name,


    # Print whole tree by calling print_directed_trees()
    # for each tree in tree_dict
    def print_whole_tree(self):
        for key in tree_dict:
            tree_dict[key].print_directed_trees()
            print ""
        return

    # Sets the self.visited values all to False using tree_dict
    def set_tree_false(self):
        for key in tree_dict:
            tree_dict[key].visited = False
        return

# dfs() is the depth first search called by a user
# and takes the current tree, and the tree to search for
# it first sets all the self.visited flags to False
# it then calls depth() which is the recursive search function
# it returns either the tree of a successful search or None
# if nothing found
def dfs(tree, search_name):
    # set all the self.visited values to False
    tree.set_tree_false()
    # create a new tree value to hold what depth() returns
    new_tree = depth(tree, search_name)
    #return new_tree
    return new_tree


# depth is the recursive function called by dfs() 
# to find if the starting tree points to the searched tree
# otherwise, if no matching value is found, None is returned
def depth(tree, search_name):
    # current tree is marked as visited
    tree.visited = True

    # if this tree matches the search, return the tree
    if tree.name == search_name:
        return tree 

    # For each tree pointed to, recursively call depth()
    for new_tree in tree.points_to:

            # Only if the new_tree has not yet been visited, send it to depth() 
            if new_tree.visited == False:
                is_match = depth(new_tree, search_name)
                
                # if a match was found, return is_match
                if is_match:
                    return is_match

    # no match was found, return None
    return None 


# search() looks through the tree_dict to find a match
# and return the tree associated with the name
# otherwise it returns None
def search(search_name):
    if tree_dict.has_key(search_name):
        return tree_dict.get(search_name)
    return None


# Return a tree based off the RC model tree
# with a couple of changes for testing:
# "h" --> "a"
# "c" --> "c"
# This tree (I think) is best thought of as a directed graph
def return_tree():
    tree=Tree("a")
    tree.add_tree("b")
    tree.add_tree("c")
 
    tree = search("b")
    tree.add_tree("d")
    tree.add_tree("e")
    tree.add_tree("f")

    tree = search("c")
    tree = tree.add_tree("g")
    tree = tree.add_tree("h")
    tree = tree.add_tree("a")

    tree = search("c")
    tree = tree.add_tree("c")
    return tree

# main function which shows a little of the testing I've done
def main():
    tree = return_tree()
    start_tree = tree
    print "current tree is ", tree.name 
    print "calling dfs() looking for a "
    tree = dfs(tree, "a")
    print "current tree is ", tree.name 
    print "calling dfs() looking for c "
    tree = dfs(tree, "c")
    print "current tree is ", tree.name 
    print "calling dfs() looking for f "
    tree = dfs(tree, "f")
    print "current tree is ", tree.name 
    print "calling dfs() looking for a "
    tree = dfs(tree, "a")

    # show that dfs() returned None, but that its still in the tree 
    if tree == None and search("a"):
        print "cant get there from here! "

    # look for something we know isn't in the tree
    tree = start_tree
    tree.print_whole_tree()
    print "lets look for z"
    print "current tree is ", tree.name 
    print "calling dfs() looking for z "
    tree = dfs(tree, "z") 
    # show that dfs() returned None and its not in the tree anyways 
    if tree == None and not search("z"):
        print "Not in this tree!" 

if __name__ == "__main__":
    main()