Submission #55097292

---

Source Code Expand

#include <iostream>
#include <vector>
using namespace std;

const int MOD = 998244353;

// モジュラ逆元を求める関数
int mod_inv(int x) {
    int result = 1, base = x, power = MOD - 2;
    while (power) {
        if (power & 1) result = 1LL * result * base % MOD;
        base = 1LL * base * base % MOD;
        power >>= 1;
    }
    return result;
}

int solve(int N, int K) {
    vector<vector<int>> dp(K + 1, vector<int>(N + 1, 0));
    dp[0][1] = 1;  // 初めは黒いボールが1番目にある

    int inv_N2 = mod_inv(N * N);
    int p_stay = (N * N - 2 * (N - 1)) * 1LL * inv_N2 % MOD;
    int p_move = 2 * 1LL * inv_N2 % MOD;

    for (int i = 1; i <= K; ++i) {
        for (int j = 1; j <= N; ++j) {
            dp[i][j] = dp[i - 1][j] * 1LL * p_stay % MOD;
            for (int k = 1; k <= N; ++k) {
                if (k != j) {
                    dp[i][j] = (dp[i][j] + dp[i - 1][k] * 1LL * p_move % MOD) % MOD;
                }
            }
        }
    }

    int expected_value = 0;
    for (int j = 1; j <= N; ++j) {
        expected_value = (expected_value + j * 1LL * dp[K][j]) % MOD;
    }

    return expected_value;
}

int main() {
    int N, K;
    cin >> N >> K;
    int result = solve(N, K);
    cout << result << endl;
    return 0;
}

Submission Info

Judge Result