AtCoder

ABC300 E問題(Dice Product 3)を解く

AtCoder_ABC300_E

AtCoder が提供しているABC(AtCoder Beginner Contest)300 のE問題をC++とPythonで解いてみました。ABC300は、2023年4月29日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

E問題 Dice Product 3(Difficulty : 1354)

問題はリンク先をご覧ください。

ABC300 E問題 Dice Product 3

確率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問題解説)が理解できたため、この問題も公式解説を読みながらプログラムを書きました。

すこしずつ理解を深めていこうと思います。

引き続き ABC に挑戦して、解説記事を書いていくつもりです。

COMMENT

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

CAPTCHA