提出 #36447762

---

ソースコード 拡げる

#include <bits/stdc++.h>
#include <atcoder/modint>

#ifdef PK
#include <debug.hpp>
#else
#define debug(...)
#endif
 
using namespace std;
 
#define all(a) (a).begin(), (a).end()
#define int long long

using mint = atcoder::modint998244353;
// using mint = double;

int32_t main() {
    cin.tie(0); ios_base::sync_with_stdio(false);
 
    // E[visited blue], E[(visited blue)^2] for each
    // vertex after i stages
    // for every adjacent, 
    // if it is blue, add 1/deg * (1 + E[visited blue]) otherwise 1/deg * E[visited blue]
    // if it is blue, add 1/deg * (E[(visited blue)^2] + 2 * E[visited blue] + 1) otherwise 1/deg * E[(visited blue)^2]
    int n, m, k;
    cin >> n >> m >> k;
    vector<vector<int>> adj(n);
    for (int i = 0; i < m; i++) {
        int a, b;
        cin >> a >> b;
        --a, --b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    vector<int> c(n);
    for (auto &i: c) cin >> i;
    vector<mint> prob(n), dp(n), dp2(n);
    prob[0] = 1;
    mint res = 0;
    for (int mv = 0; mv < k; mv++) {
        vector<mint> nprob(n), ndp(n), ndp2(n);
        for (int at = 0; at < n; at++) {
            int deg = adj[at].size();
            auto lol = prob[at] / deg;
            for (auto&nxt: adj[at]) {
                nprob[nxt] += lol;
                ndp[nxt] += lol * ((1 - c[nxt]) + dp[at]);
                ndp2[nxt] += lol * (dp2[at] + (1 - c[nxt]) * (2 * dp[at] + 1));
            }
        }
        for (int nxt = 0; nxt < n; nxt++) {
            if (ndp[nxt].val() == 0 && ndp2[nxt].val() == 0)  
                continue;
            ndp[nxt] /= nprob[nxt];
            ndp2[nxt] /= nprob[nxt];
        }
        dp = ndp, dp2 = ndp2, prob = nprob;
        for (int i = 0; i < n; i++) {
            if (c[i]) res += dp2[i] * prob[i];
        }
    }
    cout << res.val() << '\n';
    // cout << res << '\n';
}

提出情報

ジャッジ結果