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