けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 211 D - Number of Shortest paths (茶色, 400 点)

「最適解の個数」を求めるという超典型!

問題概要

頂点数  N、辺数  M の単純無向グラフが与えられます。

頂点  1 から頂点  N へと至る最短経路の本数を 1000000007 で割った余りを答えてください。

制約

  •  2 \le N \le 2 \times 10^{5}
  •  0 \le M \le 2 \times 10^{5}

解法

最短路を求めるだけならば、BFS を使うことで求められます。

そして、最短路の個数も一緒に数え上げる方法については、以前にも次の記事で解説しています。

drken1215.hatenablog.com

通常の BFS では、辺  (u, v) を用いて頂点  v へと至る最短路長 dp[v] を更新する際には

if (dp[v] == -1) {
    que.push(v);
    dp[v] = dp[u] + 1;
}

といった処理をします。最短路の個数 cnt[v] も同時に数えたい場合には、次のようにすればよいでしょう。

if (dp[v] == -1) {
    que.push(v);
    dp[v] = dp[u] + 1;
    cnt[v] = cnt[u];  // 新たに更新するため
} else if (dp[v] == dp[u] + 1) {
    cnt[v] += cnt[u];  // 新たに増える
}

たったこれだけの変更でこの問題は解けます。計算量は通常の BFS と同じく  O(N + M) です。

コード

#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using Graph = vector<vector<int>>;
using mint = atcoder::modint1000000007;

int main() {
    int N, M;
    cin >> N >> M;
    Graph G(N);
    for (int i = 0; i < M; ++i) {
        int A, B;
        cin >> A >> B;
        --A, --B;
        G[A].push_back(B);
        G[B].push_back(A);
    }
    
    queue<int> que;
    vector<int> dp(N, -1);
    vector<mint> cnt(N, 0);
    que.push(0);
    dp[0] = 0;
    cnt[0] = 1;
    while (!que.empty()) {
        int u = que.front();
        que.pop();
        for (auto v : G[u]) {
            if (dp[v] == -1) {
                // 新たに更新される場合
                que.push(v);
                dp[v] = dp[u] + 1;
                cnt[v] = cnt[u];

            } else if (dp[v] == dp[u] + 1) {
                // 最短路が増える場合
                cnt[v] += cnt[u];
            }
        }
    }
    cout << cnt[N-1].val() << endl;
}