Search code examples
pythonrecursiondepth-first-searchtree-search

Pausing a recursive depth-first search


I have a basic recursive depth first search program that searches unexplored paths in a large tree.
Is there a way to pause and save the tree so I can resume searching later rather than restarting the search every time I end my program?
I am using python and a python IDE called pycharm.


Solution

  • You would have to use an iterative version of DFS by keeping a stack object in place of the inbuilt recursion stack.

    Then we can listen to a Ctrl+C button press event to signal that the DFS should be paused.
    When its paused, we can pickle(serialize) the dfs state to a file.
    When the script is restarted, it will detect the presence of this pickle file and will resume the DFS.

    import pickle
    import os
    import multiprocessing
    import signal
    import time
    
    
    ctrl_c_pressed_event = multiprocessing.Event()
    ctrl_c_pressed_event.clear()
    
    def ctrl_c_handler(*args):
        ctrl_c_pressed_event.set()
    
    signal.signal(signal.SIGINT, ctrl_c_handler)
    
    GRAPH_PICKLE_PATH = 'graph_pickle.pkl'
    
    
    class DfsState:
        def __init__(self, visited=set(), path=[], dfs_q=[]):
            self.visited = visited
            self.path = path
            self.dfs_q = dfs_q
        def __str__(self):
            return 'visited={}\npath={}\nq={}'.format(self.visited, self.path, self.dfs_q)
    
    
    class Graph:
        def __init__(self):
            # Graph is stored in a dictionary of <vertex> vs <set of vertices>
            # Each vertex is just an int
            self.neighbours = {}
            self.vertices = set()
            self.dfs_state = DfsState()
    
        def add_edge(self, src, dst):
            """
            Add an edge from `src` to `dst`.
            """
            self.vertices.add(src)
            self.vertices.add(dst)
            src_neighbours = self.neighbours.get(src, set())
            src_neighbours.add(dst)
            self.neighbours[src] = src_neighbours
    
        def get_all_edges(self):
            return self.neighbours
    
        def get_neighbours(self, src):
            return self.neighbours.get(src, set())
        
        def add_vertex(self, src):
            if src not in self.neighbours:
                self.neighbours[src] = set()
            self.vertices.add(src)
    
    
    def dfs_iterative(graph):
        visited = graph.dfs_state.visited
        path = graph.dfs_state.path
        dfs_q = graph.dfs_state.dfs_q
        while dfs_q:
            if ctrl_c_pressed_event.is_set():
                return 'PAUSED'
            current = dfs_q.pop()
            print(current)
            if current not in visited:
                visited.add(current)
                path.append(current)
                neighbours = graph.get_neighbours(current)
                for neighbour in neighbours:
                    if neighbour not in visited:
                        dfs_q.append(neighbour)
                # Just for testing
                time.sleep(1)
        return 'DONE'
    
    
    def init_graph():
        graph = Graph()
        graph.add_edge(1, 2)
        graph.add_edge(1, 3)
        graph.add_edge(2, 3)
        graph.add_edge(2, 5)
        graph.add_edge(3, 2)
        graph.add_edge(3, 4)
        graph.add_edge(4, 5)
        graph.add_edge(5, 6)
        graph.add_edge(5, 7)
        graph.add_edge(7, 8)
        start_vertex = 1
        graph.dfs_state.dfs_q = [start_vertex]
        return graph
    
    
    def pickle_graph(graph):
        pickle.dump(graph, open(GRAPH_PICKLE_PATH, 'wb'))
    
    
    def load_pickle_graph():
        return pickle.load(open(GRAPH_PICKLE_PATH, "rb" ))
    
    
    def main():
        if os.path.exists(GRAPH_PICKLE_PATH):
            graph = load_pickle_graph()
            print('Resuming DFS')
            print(graph.dfs_state)
        else:
            graph = init_graph()
        status = dfs_iterative(graph)
        pickle_graph(graph)
        print('\n' + str(graph.dfs_state))
        if status == 'PAUSED':
            print('DFS paused!')
        else:
            print('DFS complete!')
    
    
    if __name__ == "__main__":
        main()
    

    Output:

    >>> ajaggi-mn1:~/Documents/personal/practice/python$ python a.py
    1
    3
    4
    ^C
    visited=set([1, 3, 4])
    path=[1, 3, 4]
    q=[2, 2, 5]
    DFS paused!
    
    
    >>> ajaggi-mn1:~/Documents/personal/practice/python$ python a.py
    Resuming DFS
    visited=set([1, 3, 4])
    path=[1, 3, 4]
    q=[2, 2, 5]
    5
    7
    8
    ^C
    visited=set([1, 3, 4, 5, 7, 8])
    path=[1, 3, 4, 5, 7, 8]
    q=[2, 2, 6]
    DFS paused!
    
    
    >>> ajaggi-mn1:~/Documents/personal/practice/python$ python a.py
    Resuming DFS
    visited=set([1, 3, 4, 5, 7, 8])
    path=[1, 3, 4, 5, 7, 8]
    q=[2, 2, 6]
    6
    2
    2
    
    visited=set([1, 2, 3, 4, 5, 6, 7, 8])
    path=[1, 3, 4, 5, 7, 8, 6, 2]
    q=[]
    DFS complete!
    
    
    >>> ajaggi-mn1:~/Documents/personal/practice/python$ python a.py
    Resuming DFS
    visited=set([1, 2, 3, 4, 5, 6, 7, 8])
    path=[1, 3, 4, 5, 7, 8, 6, 2]
    q=[]
    
    visited=set([1, 2, 3, 4, 5, 6, 7, 8])
    path=[1, 3, 4, 5, 7, 8, 6, 2]
    q=[]
    DFS complete!