AtCoder-ABC204 C - Tour【Python解答例】

AtCoder Beginner Contest204のC問題についてPythonの解答例を記事にしていきます。
AtCoder Beginner Contest 204 - AtCoder

AtCoder Beginner Contest204 C - Tour

C - Tour

問題文

AtCoder 国には 1 から N の番号がついた N 個の都市と、1 から M の番号がついた M 個の道路があります。

道路 i を通ると都市 Ai から Bi へ移動することができます。都市 Bi から都市 Ai への通行はできません。

ピューマは、どこかの都市からスタートし、0 本以上の道路を使い移動して、どこかの都市をゴールとするような旅行の計画を立てようとしています。

スタート地点とゴール地点の都市の組として考えられるものは何通りありますか？

制約

・2≤N≤2000
・0≤M≤min(2000,N(N−1))
・1≤Ai,Bi≤N
・Ai≠Bi
・(Ai,Bi) は相異なる
・入力に含まれる値は全て整数である

解答例

import collections

n, m = map(int,input().split())

Graph = [[] for _ in range(n)]
for i in range(m):
    a, b = map(int,input().split())
    a -= 1
    b -= 1
    Graph[a].append(b)

# i 番目の都市からたどり着ける都市を調べるbfs
def bfs(i):
    global ans
    # 未到達で0、到達で1
    visited = [0] * n

    que = collections.deque()
    que.append(i)
    visited[i] = 1
    ans += 1

    while len(que) != 0:
        tar = que.popleft()
        for j in Graph[tar]:
            # 道でつながって都市が未到達ならキューに追加
            if visited[j] == 0:
                que.append(j)
                visited[j] = 1
                ans += 1

    return


ans = 0
for i in range(n):
    bfs(i)

print(ans)

解説

N個の都市とM個の片道通行の道があり、スタート地点とゴール地点の都市が何通りあるか答える問題です。

i番目の都市をスタート地点としたときに、たどり着くことのできる（ゴール地点にできる）都市の数を幅優先探索BFSで探索しました。

BFS関数ではvisited配列を作成し、到達した都市の要素は1としていきました。
到達した都市をpopして、道でつながっている都市を調べていきます。
つながっている都市が訪れたことがなければ（visitedが0なら）、その都市の番号をスタックし、visitedを1に更新します。
求める組み合わせであるans変数には1を足していきました。

このBFSをすべての都市に対して行うことで、ansに求める数が格納されます。

最後にansを出力すればＡＣでした！

ABC204の関連記事はこちら
ebisuke33.hatenablog.com