AtCoder が提供しているABC(AtCoder Beginner Contest)298 のE問題をC++とPythonで解いてみました。ABC298は、2023年4月15日21:00に実施されました。
この回は、外部からDDoS攻撃を受けてサーバが不安定になり unrated となりました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Unfair Sugoroku(Difficulty : 1341)
問題はリンク先をご覧ください。
確率DPです。ゴールから値を決めていきます。AtCoder Problems による Difficulty は 1341 でした。
解答案
確率DPとして解いていきます。dp[i][j][k] は、高橋君の位置を i、青木君の位置を j、k が0のときは高橋君の手番、1のときは青木君の手番とし、高橋君の勝つ確率を表します。明らかに以下が成り立ちます(手番に関係なく高橋君が勝っています)。
- dp[N][j][0] = 1 (j は 1 から N-1)
- dp[N][j][1] = 1 (j は 1 から N-1)
また以下も成立します(手番に関係なく青木君が勝っています)。
- dp[i][N][0] = 0 (i は 1 から N-1)
- dp[i][N][1] = 0 (i は 1 から N-1)
確率を表す漸化式は、以下となります。高橋君の場合(K=0の場合)は、1からPまで場所が進みますが、手番が入れ替わるため k の値が反転します。また1からPまでは同じ確率となるため、加えてからPで割っています。青木君の場合も同様です。
- dp[i][j][0] = 1 / P (dp[i+1][j][1] + dp[i+2][j][1] + … + dp[i+P][j][1]
- dp[i][j][1] = 1 / Q (dp[i][j+1][1] + dp[i][j+2][1] + … + dp[i][j+Q][1]
C++ プログラム例(ABC298E)
考察した式をもとにプログラムを実装します。プログラムの補足です。
- すごろくはNを超えないため、Nを超えた場合はNとする。
- PとQで割るために MOD-2 乗を求めています(フェルマーの小定理)。
以下が、C++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
#define MOD 998244353
typedef unsigned long long int ull;
ull power(ull n, ull p, ull mod)
{
if (p == 0) {
return 1;
}
if (p % 2 == 0) {
ull t = power(n, p / 2, mod);
return t * t % mod;
}
return (n * power(n, p - 1, mod))% mod;
}
int main()
{
int n, a, b, p, q;
cin >> n >> a >> b >> p >> q;
vector<vector<vector<ull>>> dp(n + 1, vector<vector<ull>>(n + 1, vector<ull>(2, 0)));
for (int i = 0; i < n; ++i) {
dp[n][i][0] = 1;
dp[n][i][1] = 1;
dp[i][n][0] = 0;
dp[i][n][1] = 0;
}
ull invP = power(p, MOD - 2, MOD);
ull invQ = power(q, MOD - 2, MOD);
for (int i = n - 1; i >= 0; --i) {
for (int j = n - 1; j >= 0; --j) {
for (int k = 1; k <= p; ++k) {
dp[i][j][0] += dp[min(n, i + k)][j][1];
}
dp[i][j][0] *= invP;
dp[i][j][0] %= MOD;
for (int k = 1; k <= q; ++k) {
dp[i][j][1] += dp[i][min(n, j + k)][0];
}
dp[i][j][1] *= invQ;
dp[i][j][1] %= MOD;
}
}
cout << dp[a][b][0] << endl;
return 0;
}
AtCoder Library(ACL)を使ったプログラムも紹介します。ACL は、AtCoder のコンテストで使えるライブラリです。ACL が提供している機能のひとつである modint を使うことによって、自動で mod を計算することができます。
#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
typedef modint998244353 mint;
int main()
{
int n, a, b, p, q;
cin >> n >> a >> b >> p >> q;
vector<vector<vector<mint>>> dp(n + 1, vector<vector<mint>>(n + 1, vector<mint>(2, 0)));
for (int i = 0; i < n; ++i) {
dp[n][i][0] = 1;
dp[n][i][1] = 1;
dp[i][n][0] = 0;
dp[i][n][1] = 0;
}
for (int i = n - 1; i >= 0; --i) {
for (int j = n - 1; j >= 0; --j) {
for (int k = 1; k <= p; ++k) {
dp[i][j][0] += dp[min(n, i + k)][j][1];
}
dp[i][j][0] /= p;
for (int k = 1; k <= q; ++k) {
dp[i][j][1] += dp[i][min(n, j + k)][0];
}
dp[i][j][1] /= q;
}
}
cout << dp[a][b][0].val() << endl;
return 0;
}
どちらも AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC298E)
Python では、mod 付きの階乗は、pow で求めることができます(16、20行目)。以下となります。
"""AtCoder Beginner Contest 298 E"""
MOD = 998244353
n, a, b, p, q = map(int, input().split())
dp = [[[0 for _ in range(2)] for _ in range(n + 1)] for _ in range(n + 1)]
for i in range(n):
dp[n][i][0] = 1
dp[n][i][1] = 1
dp[i][n][0] = 0
dp[i][n][1] = 0
for i in range(n - 1, -1, -1):
for j in range(n - 1, -1, -1):
for k in range(1, p + 1):
dp[i][j][0] += dp[min(n, i + k)][j][1]
dp[i][j][0] *= pow(p, MOD-2, MOD)
dp[i][j][0] %= MOD
for k in range(1, q + 1):
dp[i][j][1] += dp[i][min(n, j + k)][0]
dp[i][j][1] *= pow(q, MOD-2, MOD)
dp[i][j][1] %= MOD
print(dp[a][b][0])
こちらも「AC」と判定されました。
最後に
この問題は、漸化式が頭の中で納得できるまでに時間がかかりました。DPで確率(または期待値)を処理する問題は一定の頻度ででるため、手の内化できればと思います。
引き続き ABC の問題を紹介していきます。