Implementing DFS and BFS for binary tree

123 picture 123 · Apr 24, 2016 · Viewed 11.2k times · Source

I'm trying to traverse a binary tree using depth first traversal and breadth first traversal, but I'm running into trouble. My node and tree implementation seems to be fine, I'm just not sure how to properly traverse the tree depth-wise and breadth-wise.

class Node:
    def __init__(self, val):
        self.l = None
        self.r = None
        self.v = val

class Tree:
    def __init__(self):
        self.root = None

    def getRoot(self):
        return self.root

    def add(self, val):
        if(self.root == None):
            self.root = Node(val)
        else:
            self._add(val, self.root)

    def _add(self, val, node):
        if(val < node.v):
            if(node.l != None):
                self._add(val, node.l)
            else:
                node.l = Node(val)
        else:
            if(node.r != None):
                self._add(val, node.r)
            else:
                node.r = Node(val)

    def find(self, val):
        if(self.root != None):
            return self._find(val, self.root)
        else:
            return None

    def _find(self, val, node):
        if(val == node.v):
            return node
        elif(val < node.v and node.l != None):
            self._find(val, node.l)
        elif(val > node.v and node.r != None):
            self._find(val, node.r)

    def printTree(self):
        if(self.root != None):
            self._printTree(self.root)

    def _printTree(self, node):
        if(node != None):
            self._printTree(node.l)
            print(str(node.v) + ' ')
            self._printTree(node.r)

    # This doesn't work - graph is not subscriptable
    def dfs(self, graph, start):
        visited, stack = set(), [start]
        while stack:
            vertex = stack.pop()
            if vertex not in visited:
                visited.add(vertex)
                stack.extend(graph[vertex] - visited)
        return visited

     # Haven't tried BFS.  Would use a queue, but unsure of the details.

Answer

schwobaseggl picture schwobaseggl · Apr 24, 2016

If it is a tree, visited can be a list since trees are non-circular, so there no need to check whether you have visited a node before and, more importantly, you want to maintain the order of your traversal.

def dfs(self, tree):
    if tree.root is None:
        return []
    visited, stack = [], [tree.root]
    while stack:
        node = stack.pop()
        visited.append(node)
        stack.extend(filter(None, [node.r, node.l]))  
        # append right first, so left will be popped first
    return visited