【Python】幅優先探索のいろいろな実装方法

芝の上のノートパソコン プログラミング

幅優先探索のいろいろな実装方法について解説します。

幅優先探索についての基本については、以下の記事で紹介しています。興味のある方はぜひ読んでみてください。

幅優先探索って何?Pythonで学ぶ初心者向け入門ガイド

幅優先探索の基本

幅優先探索のアルゴリズムは次のようになります。

  1. 探索開始地点をキューに追加する。
  2. キューが空になるまで次の操作を行う。
    • キューから先頭の要素を取り出す。
    • 先頭の要素から行くことのできる、未探索の地点をキューに入れて探索済みにする

コードにすると次にようになります。

# キューに探索開始地点を追加し、探索済みに
queue.append(0)
searched[0] = True

# キューが空になるまで繰り返す
while queue:
    # キューから先頭の要素を取り出す
    v = queue.popleft()

    # 先頭の要素から行ける要素について処理
    for next_v in graph[v]:

        # キューに追加しない条件
        if searched[next_v]:
            continue

        # キューに追加し、探索済みに
        queue.append(next_v)
        searched[next_v] = True

スタート地点からゴール地点にたどり着くか

AtCoder Typical Contest 001 – A

A – 深さ優先探索

スタート地点からゴール地点にたどり着くかを判定する問題です。

解答コードは次のコードです。

# collections モジュールから deque クラスをインポートする
from collections import deque

# 入力
H, W = map(int, input().split())
c = [input() for _ in range(H)]


# スタート地点とゴール地点の取得
for y in range(H):
    for x in range(W):
        # スタート地点だったとき
        if c[y][x] == 's':
            sy, sx = y, x
        
        # ゴール地点だったとき
        if c[y][x] == 'g':
            gy, gx = y, x


# キューの宣言
queue = deque()
# すべてのマスを未探索に
visited = [[False] * W for _ in range(H)]


# 探索用の変数
dx = [-1, 0, 1, 0]
dy = [0, -1, 0, 1]


# キューに追加しない条件にあたるか判定する関数
def is_skip(next_y, next_x):
    # 移動できるマスが盤面外のとき
    if (next_y < 0 or H <= next_y) or (next_x < 0 or W <= next_x):
        return True

    # 移動できるマスが壁のとき
    if c[next_y][next_x] == '#':
         return True
    
    # 移動できるマスが探索済みのとき
    if visited[next_y][next_x]:
         return True
    
    return False

  
# 幅優先探索
def bfs(sy, sx):
    # スタート地点をキューに追加する
    queue.append([sy, sx])
    # スタート地点を探索済みにする
    visited[sy][sx] = True
    
    # キューが空になるまで繰り返す
    while queue:
        # キューの先頭の要素を取り出す
        y, x = queue.popleft()

        # 先頭の要素から上下左右への処理
        for i in range(4):
            # 現在地点から移動できるマス
            next_y = y + dy[i]
            next_x = x + dx[i]

            # 移動できるマスをキューに追加しないとき
            if is_skip(next_y, next_x):
                continue
                
            # 移動できるマスをキューに追加する
            queue.append([next_y, next_x])
            # 移動できるマスを探索済みにする
            visited[next_y][next_x] = True


# スタート地点から幅優先探索
bfs(sy, sx)

# 答えの出力
if visited[gy][gx]:
    print('Yes')
else:
    print('No')

スタート地点から幅優先探索をして、ゴール地点の visited が True になっていればたどり着けるということになります。

現在のマスから上下左右の探索は、探索用のリストを準備しておけば、簡単に探索できます。

dx と dy のリストを作り、それを使うことで上下左右の探索を行っています。

# 現在のマス。(1, 1)とする
x = 1
y = 1

# 4方向移動用のリスト
dx = [-1, 0, 1, 0]
dy = [0, -1, 0, 1]

# 上下左右の探索
for i in range(4):
    next_x = x + dx[i]
    next_y = y + dy[i]
    print(next_x, next_y)
  
"""実行結果
0 1  # 上
1 0  # 左
2 1  # 下
1 2  # 右
"""

最短距離を求める

AtCoder Typical Contest 002 – A

A – 幅優先探索

スタートからゴールまでの最短距離を求める問題です。

解答コードは次のコードです。

# collections モジュールから deque クラスをインポートする
from collections import deque

# 入力
R, C = map(int, input().split())
sy, sx = map(int, input().split())
gy, gx = map(int, input().split())
c = [input() for _ in range(R)]


# スタート地点とゴール地点は管理しやすいように-1
sy -= 1
sx -= 1
gy -= 1
gx -= 1


# キューの宣言
queue = deque()
# すべてのマスを -1(未探索)で初期化
distance = [[-1] * C for _ in range(R)]


# 探索用の変数
dx = [-1, 0, 1, 0]
dy = [0, -1, 0, 1]


# キューに追加しない条件にあたるか判定する関数
def is_skip(next_y, next_x):
    # 移動できるマスが盤面外のとき
    if (next_y < 0 or R <= next_y) or (next_x < 0 or C <= next_x):
        return True

    # 移動できるマスが探索済みのとき
    if distance[next_y][next_x] != -1:
         return True

    # 移動できるマスが壁のとき
    if c[next_y][next_x] == '#':
         return True
      
    return False

  
# 幅優先探索
def bfs(sy, sx):
    # スタート地点をキューに追加する
    queue.append([sy, sx])
    # スタート地点の距離を0にする
    distance[sy][sx] = 0
    
    # キューが空になるまで繰り返す
    while queue:
        # キューの先頭の要素を取り出す
        y, x = queue.popleft()

        # 先頭の要素から上下左右への処理
        for i in range(4):
            # 現在地点から移動できるマス
            next_y = y + dy[i]
            next_x = x + dx[i]

            # 移動できるマスをキューに追加しないとき
            if is_skip(next_y, next_x):
                continue
                
            # 移動できるマスをキューに追加する
            queue.append([next_y, next_x])
            # 移動できるマスの距離を現在のマスの距離+1にする
            distance[next_y][next_x] = distance[y][x] + 1  


# スタート地点から幅優先探索
bfs(sy, sx)

# 答えの出力
print(distance[gy][gx])

distance で「探索済みか」と「スタート地点からの距離」を管理しています。

-1 で初期化しているので、distance の値が -1 であれば、未探索という確認もできます。

最大距離を求める

AtCoder Beginner Contest 151 – D

D – Maze Master

すべてのマスから幅優先探索をして、その中の最大値を出力します。

今までは、1つの地点から1回だけ幅優先探索をする問題でしたが、今回は、すべてのマスから幅優先探索をします。

解答コードは次のコードです。

# collections モジュールから deque クラスをインポートする
from collections import deque

# 入力
H, W = map(int, input().split())
S = [input() for _ in range(H)]


# 探索用の変数
dx = [-1, 0, 1, 0]
dy = [0, -1, 0, 1]


# 最大距離を返す関数
def max_distance(distance):
    max_dist = 0
    
    for dist in distance:
        max_dist = max(max_dist, max(dist))
    
    return max_dist

  
# キューに追加しない条件にあたるか判定する関数
def is_skip(next_y, next_x):
    # 移動できるマスが盤面外のとき
    if (next_y < 0 or H <= next_y) or (next_x < 0 or W <= next_x):
        return True

    # 移動できるマスが壁のとき
    if S[next_y][next_x] == '#':
         return True
    
    # 移動できるマスが探索済みのとき
    if distance[next_y][next_x] != -1:
         return True
    
    return False

  
# 幅優先探索
def bfs(sy, sx):
    # スタート地点をキューに追加する
    queue.append([sy, sx])
    # スタート地点を探索済みにする
    distance[sy][sx] = 0
    
    # キューが空になるまで繰り返す
    while queue:
        # キューの先頭の要素を取り出す
        y, x = queue.popleft()

        # 先頭の要素から上下左右への処理
        for i in range(4):
            # 現在地点から移動できるマス
            next_y = y + dy[i]
            next_x = x + dx[i]

            # 移動できるマスをキューに追加しないとき
            if is_skip(next_y, next_x):
                continue
                
            # 移動できるマスをキューに追加する
            queue.append([next_y, next_x])
            # 移動できるマスを探索済みにする
            distance[next_y][next_x] = distance[y][x] + 1

    # 探索終了後の distance の最大距離を返す
    return max_distance(distance)

  
# 答えの変数
answer = 0
  
# すべてのマスから幅優先探索
for y in range(H):
    for x in range(W):
       # マスが壁のとき
        if S[y][x] == '#':
            continue
            
        # キューの宣言
        queue = deque()
        # すべてのマスを未探索に
        distance = [[-1] * W for _ in range(H)]
        
        answer = max(answer, bfs(y, x))


# 答えの出力
print(answer)

開始地点を全探索して、幅優先探索をしてその最大距離を max() で求めています。

特定の距離の数を求める

AtCoder Beginner Contest 016 – C

C – 友達の友達

スタート地点から距離が2離れた地点の数を求める問題です。

解答コードは次のコードです。

# collections モジュールから deque クラスをインポートする
from collections import deque

# 入力
N, M = map(int, input().split())

# グラフの入力
graph = [[] for _ in range(N)]
for _ in range(M):
    A, B = map(int, input().split())
    A -= 1
    B -= 1
    graph[A].append(B) 
    graph[B].append(A)

  
# 幅優先探索
def bfs(user):
    # ユーザーをキューに追加する
    queue.append(user)
    # ユーザーを探索済みにする
    distance[user] = 0
    
    # キューが空になるまで繰り返す
    while queue:
        # キューの先頭の要素を取り出す
        friend = queue.popleft()

        # 先頭の要素から隣接要素への処理
        for next_friend in graph[friend]:

            # 次の友だちをキューに追加しないとき
            if distance[next_friend] != -1:
                continue
                
            # 次の友だちをキューに追加する
            queue.append(next_friend)
            # 次の友だちを探索済みにする
            distance[next_friend] = distance[friend] + 1

    # 探索終了後の distance が2の数を返す
    return distance.count(2)

  
# すべてのユーザーから幅優先探索
for user in range(N):
    # キューの宣言
    queue = deque()
    # すべてのユーザーを未探索に
    distance = [-1] * N
    
    # 答えの出力    
    print(bfs(user))

すべてのユーザーについて幅優先探索をして、距離が2離れた数を出力しています。

連結成分の個数を求める

AtCoder Beginner Contest 284 – C

C – Count Connected Components

連結成分の個数を求める問題です。

解答コードは次のコードです。

# collections モジュールから deque クラスをインポートする
from collections import deque

# 入力
N, M = map(int, input().split())

# グラフの入力
graph = [[] for _ in range(N)]
for _ in range(M):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    graph[u].append(v) 
    graph[v].append(u)

  
# 幅優先探索
def bfs(start):
    # 探索開始する頂点をキューに追加する
    queue.append(start)
    # 探索開始する頂点を探索済みに
    visited[start] = True
    
    # キューが空になるまで繰り返す
    while queue:
        # キューの先頭の要素を取り出す
        v = queue.popleft()

        # 先頭の要素から隣接要素への処理
        for next_v in graph[v]:

            # 次の頂点をキューに追加しないとき
            if visited[next_v]:
                continue
                
            # 次の頂点をキューに追加する
            queue.append(next_v)
            # 次の頂点を探索済みにする
            visited[next_v] = True

            
# キューの宣言
queue = deque()
# すべての頂点を未探索に
visited = [False] * N

# 答えの変数
answer = 0

# すべての頂点から幅優先探索
for v in range(N):
    # 探索済みのときはスキップ
    if visited[v]:
        continue
    
    # 幅優先探索
    bfs(v)
    
    # 連結成分の個数
    answer += 1
    
    
# 答えの出力    
print(answer)

幅優先探索をすれば、その頂点からいけるすべて頂点が探索済みになるため、幅優先探索をした回数が連結成分の個数になります。

最短距離でたどりつくの経路の数を求める

AtCoder Beginner Contest 211 – D

D – Number of Shortest paths

最短距離でたどり着く経路の数を求める問題です。

解答コードは次のコードです。

# collections モジュールから deque クラスをインポートする
from collections import deque

# 入力
N, M = map(int, input().split())

# グラフの入力
graph = [[] for _ in range(N)]
for _ in range(M):
    A, B = map(int, input().split())
    A -= 1
    B -= 1
    graph[A].append(B) 
    graph[B].append(A)

    
# キューの宣言
queue = deque()
# すべての頂点を -1(未探索)で初期化
distance = [-1] * N
# 頂点1から最短距離でたどりつく経路の数
count = [0] * N


# 幅優先探索
def bfs(start):
    # 探索開始する頂点をキューに追加する
    queue.append(start)
    # 探索開始する頂点を探索済みに
    distance[start] = 0
    # 探索開始地点への経路の数を1に
    count[start] = 1
    
    # キューが空になるまで繰り返す
    while queue:
        # キューの先頭の要素を取り出す
        city = queue.popleft()

        # 先頭の要素から隣接要素への処理
        for next_city in graph[city]:

            # 次の頂点が未探索のとき
            if distance[next_city] == -1:
                # 次の頂点をキューに追加する
                queue.append(next_city)
                # 次の頂点を探索済みにする
                distance[next_city] = distance[city] + 1
                # 次の頂点への最短経路の数
                count[next_city] = count[city]
                continue
            
            # 探索済みで最短距離のとき
            if distance[next_city] == distance[city] + 1:
                count[next_city] += count[city]
                count[next_city] %= 10**9 + 7


# 幅優先探索
bfs(0)

# 答えの出力
print(count[N - 1])

count に最短距離でたどりつく経路の数を管理しています。

グラフを dict で管理する

AtCoder Beginner Contest 277 – C

C – Ladder Takahashi

幅優先探索をして、到達できる最大値を求める問題です。

グラフの実装にコツが必要です。

解答コードは次のコードです。

# collections モジュールから deque クラスをインポートする
from collections import deque

# 入力
N = int(input())

# グラフの入力
graph = dict()
for _ in range(N):
    A, B = map(int, input().split())
    # 初めてなので空のリストを作成
    if A not in graph:
        graph[A] = []
    if B not in graph:
        graph[B] = []
    graph[A].append(B) 
    graph[B].append(A)

    
# キューの宣言
queue = deque()
# 探索済みかを管理する変数
visited = set()


# 幅優先探索
def bfs(start):
    # 探索開始する頂点をキューに追加する
    queue.append(start)
    # 探索開始する頂点を探索済みに
    visited.add(start)
    
    # キューが空になるまで繰り返す
    while queue:
        # キューの先頭の要素を取り出す
        floor = queue.popleft()

        # 今いる階にはしごがないときはスキップ
        if floor not in graph:
            continue
        
        # 先頭の要素から隣接要素への処理
        for next_floor in graph[floor]:

            # 次の頂点が未探索のとき
            if next_floor in visited:
                continue
              
            # 次の頂点をキューに追加する
            queue.append(next_floor)
            # 次の頂点を探索済みにする
            visited.add(next_floor)


# 幅優先探索
bfs(1)

# 答えの出力
print(max(visited))

グラフの頂点の数をリストではなく dict で実装しています。

探索済みかどうかを管理する変数も、set を使っています。

参考

BFS (幅優先探索) 超入門! 〜 キューを鮮やかに使いこなす 〜 – Qiita

問題解決力を鍛える!アルゴリズムとデータ構造

競技プログラミングの鉄則 ~アルゴリズム力と思考力を高める77の技術~

プログラミングコンテスト攻略のためのアルゴリズムとデータ構造

プログラミングコンテストチャレンジブック [第2版] ~問題解決のアルゴリズム活用力とコーディングテクニックを鍛える~

タイトルとURLをコピーしました