AtCoder が提供しているABC(AtCoder Beginner Contest)275 のE問題をC++とPythonで解いてみました。ABC275は、2022年10月29日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Sugoroku 4(Difficulty : 1264)
問題はリンク先をご覧ください。
確率を求めるDPの典型問題でした。AtCoder Problems による Difficulty は、1264 でした。
解答案
dp[i][j] を ルーレットを i 回回したときにマス j にいる確率とします。漸化式は以下となります。
dp[i + 1][j + (1 から m)] += dp[i][j] / m
C++ プログラム例(ABC275E)
考察した漸化式を mod 998244353 で計算しました。
MOD付の演算は、足し算、引き算、掛け算は自然にできますが、割り算は工夫が必要です。フェルマーの小定理を使えば、$N \pmod M$ で割ることは、$N^{M-2} \pmod M$ を掛けることと同じになります。
この計算をするために累乗を計算する関数 power を使います(7ー17行目)。この関数は自前ライブラリとして整備したものです。$M$ で割る場合に掛ける値を変数 invM として計算しておきます(26行目)。
問題の制約から以下の2点に注意します。
- マス N にいる場合は、ルーネットを回さない(29-33行目)。
- マス N を超えた場合は、超える分だけ N から引き返す(36ー38行目)。
以下が、C++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
#define MOD 998244353
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()
{
ull n, m, k;
cin >> n >> m >> k;
vector<vector<ull>> dp(k + 1, vector<ull>(n + 1, 0));
dp[0][0] = 1;
ull invM = power(m, MOD - 2, MOD);
for (int i1 = 0; i1 < k; ++i1) {
for (int i2 = 0; i2 <= n; ++i2) {
if (i2 == n) {
dp[i1 + 1][i2] += dp[i1][i2];
dp[i1 + 1][i2] %= MOD;
continue;
}
for (int i3 = 1; i3 <= m; ++i3) {
int next = i2 + i3;
if (next > n) {
next = n - (next - n);
}
dp[i1 + 1][next] += dp[i1][i2] * invM;
dp[i1 + 1][next] %= MOD;
}
}
}
cout << dp[k][n] << endl;
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC275E)
基本的な考え方は同じです。MOD計算付きの累乗は、組込み関数 pow で計算することができます(7行目)。
"""AtCoder Beginner Contest 275 E"""
MOD = 998244353
n, m, k = map(int, input().split())
dp = [[0 for _ in range(n + 1)] for _ in range(k + 1)]
dp[0][0] = 1
invM = pow(m, MOD - 2, MOD)
for i1 in range(k):
for i2 in range(n + 1):
if i2 == n:
dp[i1 + 1][i2] += dp[i1][i2]
dp[i1 + 1][i2] %= MOD
continue
for i3 in range(1, m + 1):
nxt = i2 + i3
if nxt > n:
nxt = n - (nxt - n)
dp[i1 + 1][nxt] += dp[i1][i2] * invM
dp[i1 + 1][nxt] %= MOD
print(dp[k][n])
こちらも「AC」と判定されました。
最後に
確率を用いるDPは、ABC323E問題(解説記事)やABC326E問題(解説記事)で解いていました。この問題も確率を使うDPとしては、解きやすく感じました。
引き続き ABC の問題を紹介していきます。