[Python] BFSによるグラフの探索

無向グラフをBFSで探索するアルゴリズムを Python で記述します。

以下の続きです。

幅優先探索

幅優先探索（はばゆうせんたんさく、英: breadth first search）はグラフ理論(Graph theory)において木構造(tree structure)やグラフ(graph)の探索に用いられるアルゴリズム。アルゴリズムは根ノードで始まり隣接した全てのノードを探索する。それからこれらの最も近いノードのそれぞれに対して同様のことを繰り返して探索対象ノードをみつける。

出典: フリー百科事典『ウィキペディア（Wikipedia）』

幅優先探索は、Queue を用いて実装します。

Queue は以下を用います。

deque オブジェクト

以下のグラフを探索します。

以前のコードに下記を追加します。

• deque を使うため、1行目で collection を import する。
• visited を DFS、BFS で共用できるよう、29行目に visited フラグを False にする関数を追加し、探索を開始する際に、 visited フラグ を全て False にするようにする。
• 95行目と107行目で BFS を実装する。

import collections

class Vertex(object):
    def __init__(self, id):
        self.id = id
        self.adjacent = {}
        # DFS BFS 共用フラグ
        self.visited = False

    def add_neighbor(self, neighbor, weight=0):
        self.adjacent[neighbor] = weight
    
    def get_neighbors(self):
        return self.adjacent

    def get_connections(self):
        return self.adjacent.keys()  

    def get_vertex_id(self):
        return self.id

    def get_weight(self, neighbor):
        return self.adjacent[neighbor]

    # DFS BFS 共用フラグ
    def set_visited(self):
        self.visited = True
    
    def unset_visited(self):
        self.visited = False


class Graph(object):
    def __init__(self):
        self.vertex_dict = {}
        self.num_vertex = 0
    
    # DFS ループで扱いやすくする。
    def __iter__(self):
        return iter(self.vertex_dict.values())

    def add_vertex(self, id):
        self.num_vertex = self.num_vertex + 1
        new_vertex = Vertex(id)
        self.vertex_dict[id] = new_vertex
        return new_vertex

    def get_vertex(self, id):
        if id in self.vertex_dict:
            return self.vertex_dict[id]
        else:
            return None

    def add_edge(self, frm, to, weight=0):
        if frm not in self.vertex_dict:
            self.add_vertex(frm)
        if to not in self.vertex_dict:
            self.add_vertex(to)
        self.vertex_dict[frm].add_neighbor(self.vertex_dict[to], weight)
        # 有向グラフの場合は別
        self.vertex_dict[to].add_neighbor(self.vertex_dict[frm], weight)

    def get_vertices(self):
        return self.vertex_dict.keys()

    def get_edges(self):
        edges = []
        for v in self.vertex_dict.values():
            for w in v.get_connections():
                vid = v.get_vertex_id()
                wid = w.get_vertex_id()
                edges.append((vid, wid, v.get_weight(w)))
        return edges


def dfs(graph, current_vertex, traversal_verticies):
    current_vertex.set_visited()
    traversal_verticies.append(current_vertex.get_vertex_id())
    for neighbor in current_vertex.get_connections(): 
        if not(neighbor.visited):  
            dfs(graph, neighbor, traversal_verticies) 
 
def dfs_traversal(graph):
    # visitedフラグの初期化
    for vertex in graph:
        vertex.unset_visited()

    for current_vertex in graph:
        traversal_verticies = []
        if not(current_vertex.visited):  
            dfs(graph, current_vertex, traversal_verticies) 
        if traversal_verticies:
            print(f'DFS start from {current_vertex.get_vertex_id()}:', traversal_verticies)

def bfs(graph, start, traversal_verticies):
    queue = collections.deque()
    queue.append(start)

    while len(queue) > 0:
        current_vertex = queue.popleft()
        traversal_verticies.append(current_vertex.get_vertex_id())
        for neighbor in current_vertex.get_connections():
            if not(neighbor.visited) and (neighbor not in queue):
                queue.append(neighbor)
            current_vertex.visited = True

def bfs_traversal(graph):
    # visitedフラグの初期化
    for vertex in graph:
        vertex.unset_visited()

    for start_vertex in graph:
        traversal_verticies = []
        if not(start_vertex.visited):
            bfs(graph, start_vertex, traversal_verticies)
        if traversal_verticies:
            print(f'BFS start from {start_vertex.get_vertex_id()}:', traversal_verticies)

    
if __name__ == '__main__':

    graph = Graph()
    graph.add_vertex('a')
    graph.add_vertex('b')
    graph.add_vertex('c')
    graph.add_vertex('d')
    graph.add_vertex('e')
    graph.add_vertex('f')
    graph.add_edge('a', 'b', 1)  
    graph.add_edge('a', 'c', 1)
    graph.add_edge('a', 'e', 1)
    graph.add_edge('b', 'd', 1)
    graph.add_edge('b', 'e', 1)
    graph.add_edge('c', 'd', 1)
    graph.add_edge('c', 'e', 1)
    graph.add_edge('d', 'e', 1)

    dfs_traversal(graph)
    # start from a: ['a', 'b', 'd', 'c', 'e']
    # start from f: ['f']
    
    bfs_traversal(graph)
    # start from a: ['a', 'b', 'c', 'e', 'd']
    # start from f: ['f']

想定通りの結果になりました。