【ABC212】E - Safety Journey

atcoder.jp

解法

都市から都市への移動を繰り返すという行為が、DPで遷移するというのに酷似しているため、DPをしたくなる。
次のようなDPを考える。
dp[i][j] ← i日目に都市jにいる旅の仕方は何通りか

ただ問題文通りにK回、N個の町から通れる道を使って遷移をしようとするとO(KN²)になってしまうので間に合わない。

そこで使えない道が一つもないと仮定して、そこから使えない道を通ってしまった場合を引くことを考える。
使えない道が一つもない場合、すべての都市から今いる都市以外のすべての都市へと遷移することができる。

貰うDPで考えてみると、
dp[i + 1][j] ← $ \sum_{j=1}^{n}dp[i][j] $ - dp[i][j]
という感じでO(N)で遷移することができる。

あとは使えない道を通ってしまった場合をdp[i + 1][j]から引けばよい。

これで合計の計算量はO(K(N+M))になるので、ACできる。

実装

void solve(long long N, long long M, long long K, std::vector<long long> U, std::vector<long long> V) {
    vector<vector<long long>> G(N); // 使えない道を辺で結んだグラフ
    for (int i = 0; i < M; i++) {
        G[U[i]-1].push_back(V[i]-1);
        G[V[i]-1].push_back(U[i]-1);
    }
      
    vector<vector<mint>> dp(K+1, vector<mint>(N+1, 0));
    dp[0][0] = 1;

    REP (i, K) {
        // 使えない道がないと仮定したとき、次の都市への行き方は何通りか
        mint sm = 0;
        REP (j, N) sm += dp[i][j];

        REP (j, N) {
            dp[i + 1][j] = sm - dp[i][j]; // 移動は必ずしないといけないから
            for (auto n : G[j]) dp[i + 1][j] -= dp[i][n]; // 使えない道を使った場合を引き算
        }
    }

    c(dp[K][0]) // 出力
}