Depth-first search on a tree (directed, acyclic graph)

dfs.py

def dfs(graph, start_node, search_node):
    """
    Depth-first search for a node in a tree (directed, acyclic graph).

    Input:
        graph (dict)
        node we are starting our search with (String)
        node we are searching for (String)
    Output: whether or not the node was found in the graph (bool)
    """

    search_path = [] # search_path of nodes we've seen so far
    to_visit = [(start_node, search_path)] # Stack of nodes to visit

    while len(to_visit):
        node, search_path = to_visit.pop()

        # Add new node to search_path so far
        search_path.append(node)

        # Check if node is what we're looking for
        if node == search_node:
            print('Found node:', node)
            print('Search path:', search_path)
            return node

        for neighbor in graph[node]:
            # Add node's neighbors to visit
            to_visit.append((neighbor, search_path))


    print ('Node not found:', search_node)
    print('Search Path:', search_path)
    return None


def test():
    """
    Run basic tests on DFS.
    Tree graph (directed, acyclic) represented by a dict.
        Key: node name
        Value: set of neighbors (children since graph is assumed to be a tree)
    """
    graph = {'a': set(['b', 'c']),
             'b': set(['d', 'e', 'f']),
             'c': set(['g']),
             'd': set(),
             'e': set(),
             'f': set(),
             'g': set(['h']),
             'h': set()}

    assert(dfs(graph, 'a', 'a') == 'a') # Test element in graph, should be found
    assert(dfs(graph, 'a', 'h') == 'h') # Test element in graph, should be found
    assert(dfs(graph, 'a', 'd') == 'd') # Test element in graph, should be found

    assert(dfs(graph, 'a', 'x') == None) # Test element not in graph


test()