Thursday, May 24, 2012

How to Implement DFS and BFS in Python

#!/usr/bin/env python

class Graph(object):
    def __init__(self):
        # the key is the vertex, the value is the list of adjacent vertices
        self.adjacents = {}
        self.nedges = 0

    def num_vertices(self):
        return self.adjacents.keys()

    def num_edges(self):
        return self.nedges

    def add_edge(self, v, w):
        if v not in self.adjacents: self.adjacents[v] = []
        self.adjacents[v].append(w)
        if w not in self.adjacents: self.adjacents[w] = []
        self.adjacents[w].append(v)
        self.nedges += 1

    def adjacent(self, v):
        return self.adjacents[v]

class DepthFirstSearch(object):
    def __init__(self, graph, start):
        self.count = 0
        self.markeddict = {}
        self.edge_to = {}
        self.start = start
        self._dfs(graph, start)

    def _dfs(self, graph, vertex):
        self.markeddict[vertex] = True
        self.count += 1
        for v in graph.adjacent(vertex):
            if not self.marked(v):
                self.edge_to[v] = vertex
                self._dfs(graph, v)

    def marked(self, vertex):
        return vertex in self.markeddict

    def has_path_to(self, vertex):
        return self.marked(vertex)

    def path_to(self, vertex):
        if not self.has_path_to(vertex): return None
        path = []
        path.append(vertex)
        while vertex != self.start:
            vertex = self.edge_to[vertex]
            path.append(vertex)
        
        path.reverse()
        return path

class BreadthFirstSearch(object):
    def __init__(self, graph, start):
        self.count = 0
        self.markeddict = {}
        self.edge_to = {}
        self.start = start
        self._bfs(graph, start)

    def _bfs(self, graph, vertex):
        self.markeddict[vertex] = True
        self.count += 1
        q = []
        q.append(vertex)
        while q:
            key = q.pop(0)
            for v in graph.adjacent(key):
                if not self.marked(v):
                    q.append(v)
                    self.count += 1
                    self.markeddict[v] = True
                    self.edge_to[v] = key

    def marked(self, vertex):
        return vertex in self.markeddict

    def has_path_to(self, vertex):
        return self.marked(vertex)

    def path_to(self, vertex):
        if not self.has_path_to(vertex): return None
        path = []
        path.append(vertex)
        while vertex != self.start:
            vertex = self.edge_to[vertex]
            path.append(vertex)
        
        path.reverse()
        return path

if __name__ == "__main__":
    # key is the vertex and value is the list of adjacent vertices
    graph = Graph()
    
    graph.add_edge("0", "2")
    graph.add_edge("3", "5")
    graph.add_edge("3", "4")
    graph.add_edge("0", "1")
    graph.add_edge("1", "2")
    graph.add_edge("2", "3")
    graph.add_edge("2", "4")
    graph.add_edge("0", "5")

    dfs = DepthFirstSearch(graph, "0")
    assert "0" == "-".join(dfs.path_to("0"))
    assert "0-2-1" == "-".join(dfs.path_to("1"))
    assert "0-2" == "-".join(dfs.path_to("2"))
    assert "0-2-3" == "-".join(dfs.path_to("3"))
    assert "0-2-3-4" == "-".join(dfs.path_to("4"))
    assert "0-2-3-5" == "-".join(dfs.path_to("5"))
    assert 6 == dfs.count
    
    bfs = BreadthFirstSearch(graph, "0")
    assert "0" == "-".join(bfs.path_to("0"))
    assert "0-1" == "-".join(bfs.path_to("1"))
    assert "0-2" == "-".join(bfs.path_to("2"))
    assert "0-2-3" == "-".join(bfs.path_to("3"))
    assert "0-2-4" == "-".join(bfs.path_to("4"))
    assert "0-5" == "-".join(bfs.path_to("5"))
    assert 6 == bfs.count

No comments:

Post a Comment