"""
 breadth_depth.py 

 Looking at breadth-first and depth-first search of a graph (or tree)
 using a stack and a queue
 
$ python breadth_depth.py
 -- depth-first search
A
C
G
F
B
E
K
J
D
I
H
 -- breadth-first search
A
B
C
D
E
F
G
H
I
J
K
 
"""

# -- start graph API -------------------

a='A'; b='B'; c='C'; d='D'; e='E'; f='F'; g='G'; h='H'; i='I'; j='J'; k='K'

# each pair is (parent, child) in the tree
edges = ( (a, b),
          (a, c),
          (b, d),
          (b, e),
          (c, f),
          (c, g),
          (d, h),
          (d, i),
          (e, j),
          (e, k)
         )

def add_nodes_to_graph(graph, node1, node2):
    """ modify graph={node:set_of_neighbors} by adding node1 -> node2 """
    if not node1 in graph:
        graph[node1] = set()
    graph[node1].add(node2)

graph = {}
for (node1, node2) in edges:
    add_nodes_to_graph(graph, node1, node2)
    add_nodes_to_graph(graph, node2, node1)    

def get_neighbors(graph, node):
    """ return neighbors of a given node """
    return graph[node]
    
# -- end graph API -----------------

class Stack:
    """ first in, last out, with O(1) "x in stack"
        >>> s = Stack()
        >>> s.push(1); s.push(2); s.push(3)
        >>> 1 in s
        True
        >>> (s.pop(), s.pop(), s.pop())
        (3, 2, 1)
        >>> len(s)
        0
    """
    def __init__(self):
        self.values = []
        self.has = {}    # also put values in a dictionary
    def __len__(self):
        return len(self.values)
    def __contains__(self, value):
        # implement the "_ in _" python syntax.
        return value in self.has
    def push(self, value):
        self.values.append(value)
        self.has[value] = True
    def pop(self):
        value = self.values.pop()
        self.has.pop(value)   # remove from dictionary
        return value

class Queue(Stack):
    """ first in, first out 
        >>> q = Queue()
        >>> q.push(1); q.push(2); q.push(3)
        >>> 5 in q
        False
        >>> (q.pop(), q.pop(), q.pop())
        (1, 2, 3)
        >>> len(q)
        0
    """
    def pop(self):
        value = self.values.pop(0)
        self.has.pop(value)
        return value
    
def search(start, breadth_depth, graph, function):
    """ depth-first or breadth-first  search """
    visited = {}                 # Nodes which we're done with.
    if breadth_depth == 'depth': # Create a fringe of nodes-to-visit ...
        fringe = Stack()         # ... either a stack
    else:
        fringe = Queue()         # ... or a queue.
    fringe.push(start)           # Initialize the search.
    while len(fringe) > 0:       # Search loop:
        node = fringe.pop()                        # Get node to process.
        visited[node] = True                       # Mark it as 'processed'.
        function(node)                             # Do something with it.
        neighbors = get_neighbors(graph, node)     # Get its neighbors
        for candidate in neighbors:                # Add new ones to fringe.
            if not candidate in visited and not candidate in fringe:
                fringe.push(candidate)

                
def main():
    
    print(' -- depth-first search ')
    search(a, 'depth', graph, print)
    
    print(' -- breadth-first search ')
    search(a, 'breadth', graph, print)
    
if __name__ == '__main__':
    import doctest
    doctest.testmod()
    main()