AtCoder が提供しているABC（AtCoder Beginner Contest）300 のE問題をC++とPythonで解いてみました。ABC300は、2023年4月29日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Dice Product 3（Difficulty : 1354）
問題はリンク先をご覧ください。
確率DPで解くことができました。AtCoder Problems による Difficulty は 1354 でした。
解答案
配列 dp[i] を、持っている数が i から始めた場合に最終的に $N$ に一致する確率とします。求めるのは、dp[1] です。i を持っている状態から1回サイコロを振ると、1から6までが同じ確率ででます。このため dp[i] について以下の式が成立します。
dp[i] = 1/6(dp[i] + dp[2i] + dp[3i] + dp[4i] + dp[5i] + dp[6i])
つまり、以下となります。
5/6 dp[i] = 1/6(dp[2i] + dp[3i] + dp[4i] + dp[5i] + dp[6i])
dp[i] = 1/5(dp[2i] + dp[3i] + dp[4i] + dp[5i] + dp[6i])
dp[$N$] = 1 となります。$n > N$ について、dp[$n$] = 0 となります。上記の式を再帰的に求めます。$N$ の上限は、$10^{18}$ で約 $2^{60}, 3^{38}, 5^{26}$ となります。それほど多くの引数で再帰的に呼ばれないことが分かります。
C++ プログラム例（ABC300E）
MOD付の演算は、足し算、引き算、掛け算は自然にできますが、割り算は工夫が必要です。フェルマーの小定理を使えば、$N \pmod M$ で割ることは、$N^{M-2} \pmod M$ を掛けることと同じになります。
この計算をするために累乗を計算する関数 power を使います（11ー21行目）。この関数は自前ライブラリとして整備したものです。5で割る場合に掛ける値を変数 inv5 として計算しておきます（51行目）。
再帰的に呼び出す関数 rec で、先に考察した漸化式を計算します（37、40行目）。再帰関数の戻り値はメモ化することにより計算量を下げています（32、42行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
#define MOD 998244353
typedef unsigned long long int ull;
ull number;
ull inv5;
map<ull, ull> memo;
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;
}
ull rec(ull n)
{
if (n > number) {
return 0;
}
if (n == number) {
return 1;
}
if (memo.find(n) != memo.end()) {
return memo[n];
}
ull result = 0;
for (int i = 2; i <= 6; ++i) {
result += rec(i * n);
result %= MOD;
}
result *= inv5;
result %= MOD;
memo[n] = result;
return result;
}
int main()
{
cin >> number;
inv5 = power(5, MOD - 2, MOD);
inv5 %= MOD;
cout << rec(1) << endl;
return 0;
}AtCoder では、AtCoder Library（ACL）と呼ばれるライブラリが提供されています。
ACL で、MOD付きの計算ができる型 modint も提供されています。この問題は、modint を使うと簡潔にプログラムを書くことができます。
- 必要な宣言は、2行目、4行目、6行目です。
- modint を使うことにより漸化式がすっきり書けています（26、28行目）。
以下は、modint を使ったプログラムです。
#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
typedef modint998244353 mint;
typedef unsigned long long int ull;
ull number;
map<ull, mint> memo;
mint rec(ull n)
{
if (n > number) {
return 0;
}
if (n == number) {
return 1;
}
if (memo.find(n) != memo.end()) {
return memo[n];
}
mint result = 0;
for (int i = 2; i <= 6; ++i) {
result += rec(i * n);
}
memo[n] = result / 5;
return memo[n];
}
int main()
{
cin >> number;
cout << rec(1).val() << endl;
return 0;
}どちらも、AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC300E）
基本的な考え方は同じです。MOD計算付きの累乗は、組込み関数 pow で計算することができます（29行目）。
"""AtCoder Beginner Contest 300 E"""
def rec(n):
global MOD
global number
global inv5
global memo
if n > number:
return 0
if n == number:
return 1
if n in memo:
return memo[n]
result = 0
for i in range(2, 7):
result += rec(i * n)
result *= inv5
result %= MOD
memo[n] = result
return result
MOD = 998244353
number = int(input())
inv5 = pow(5, MOD - 2, MOD)
memo = {}
print(rec(1))こちらも「AC」と判定されました。
最後に
確率を用いるDPは苦手意識がありました。コンテスト当時のメモに「問題は読んだ。確率DPらしくあきらめる。」と書いてありました。ABC323E問題（解説）が理解できたため、この問題も公式解説を読みながらプログラムを書きました。
すこしずつ理解を深めていこうと思います。
2024/4/30追記 問題の似た目は異なりますが、ABC350E問題（解説記事）が類題となります。
引き続き ABC の問題を紹介していきます。
