rc_bfs.py

class Vertex(object):
    """represent vertices in a tree/graph."""
    def __init__(self, name):
        self.name = name
        self.children = {}
        self.status = None
        
    def add_child(self, child):
        key = child.name
        self.children[key] = child
    
    def get_children(self):
        return self.children.values()

"""
step 1: replicate tree from RC coding exercise. vertices represented with adjaceny list,
"""
A = Vertex('A')
B = Vertex('B')
C = Vertex('C')
D = Vertex('D')
E = Vertex('E')
F = Vertex('F')
G = Vertex('G')
H = Vertex('H')

A.add_child(B)
A.add_child(C)
B.add_child(D)
B.add_child(E)
B.add_child(F)
C.add_child(G)
G.add_child(H)


"""
step 2: implement BFS search
"""

def search_tree(source, target):
    """
    Searches tree for a node. To use, define source and target.
    
    For example:
    
    s = A  # root is node A
    t = 'g'  # we want to find node named 'g'
    result = search_tree(s, t)  # g if exists, else False
    
    Parameters
    ----------
    source : vertex object
    target : string
        represents the query. we will search the tree with this string.
    
    Returns
    -------
    ret : vertex object
        returns queried node if it exists in the tree, otherwise False is returned.
    """
    counter = 1
    source.status = counter
    bfs_queue = [source]
    result = False

    while len(bfs_queue) > 0:
        temp_queue = []  # stores the next layer of children

        # inspect each vertex in each layer 
        while len(bfs_queue) > 0:
            curr = bfs_queue[len(bfs_queue)-1]        
            
            # return node if it matches query
            if curr.name == target:
                return curr
            
            # otherwise, for each vertex, add children to queue
            children = curr.get_children()
            for child in children:
                counter += 1
                child.status = counter
                temp_queue.append(child)

            # remove vertex from bfs queue once inspected
            bfs_queue.pop()

        # repeat, but with next layer
        bfs_queue = temp_queue
    return result