AtCoder が提供しているABC（AtCoder Beginner Contest）360 E問題をC++とPythonで解いてみました。ABC360は、2024年6月30日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Random Swaps of Balls（Difficulty : 1249）
問題の詳細は、リンク先をご覧ください。
ABC360 E問題 Random Swaps of Balls
黒いボールが一番左にある確率を求めます。AtCoder Problems による Difficulty は 1249 でした。
解答案
コンテスト中に解くことができませんでした。公式解説を参考にコーディングしました。
C++ プログラム例（ABC360E）
公式解説の式を、そのままコーディングしました。dp[i] は、i回目の操作の後に一番左に黒がある確率です。k回の操作後に2番目以降に黒がある確率は、一律 $(1 – dp_k)/(N – 1)$ となります。期待値は、等差数列の和の公式を使って、
$$dp_k + 2 \times \frac{1 – dp_k}{N – 1} + 3 \times \frac{1 – dp_k}{N – 1} \cdots N \times \frac{1 – dp_k}{N – 1}$$
$$ = dp_k + \frac{(N + 2)(N – 1)}{2} \times \frac{1 – dp_k}{N – 1} = \frac{(N + 2)}{2} \times (1 – dp_k)$$
となります。以下が、C++プログラムです。
#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
typedef modint998244353 mint;
int main()
{
int n, k;
cin >> n >> k;
mint mint_n = (mint)n;
vector<mint> dp(k + 1);
dp[0] = mint(1);
mint p = (mint)2 * (mint_n - 1) / (mint_n * mint_n);
mint q = (mint)2 / (mint_n * mint_n);
for (int i = 0; i < k; ++i) {
dp[i + 1] = (1 - p) * dp[i] + q * (1 - dp[i]);
}
mint result = dp[k];
result += (mint_n + 2) * (1 - dp[k]) / 2;
cout << result.val() << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC360E）
Python版も基本的な考え方は同じです。以下がプログラムです。
"""AtCoder Beginner Contest 360 E"""
MOD = 998244353
n, k = map(int, input().split())
inv_n = pow(n, MOD - 2, MOD)
inv_2 = pow(2, MOD - 2, MOD)
dp = [0] * (k + 1)
dp[0] = 1
p = 2 * (n - 1) * inv_n * inv_n % MOD
q = 2 * inv_n * inv_n % MOD
for i in range(k):
dp[i + 1] = (1 + MOD - p) * dp[i] + q * (1 + MOD - dp[i])
dp[i + 1] %= MOD
result = dp[k]
result += (n + 2) * (1 + MOD - dp[k]) * inv_2
result %= MOD
print(result)こちらも「AC」と判定されました。
最後に
確率DPは、慣れてきた気がしていましたが、まだまだでした。
引き続き ABC の問題を紹介していきます。
