ABC326 E問題（Revenge of “The Salary of AtCoder Inc.”）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）326 のE問題をC++で解いてみました。ABC326は、2023年10月28日21:00に実施されました。

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

E問題 Revenge of “The Salary of AtCoder Inc.”（Difficulty : 1363）

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

ABC326 E問題 Revenge of “The Salary of AtCoder Inc.”

期待値を後ろから計算していきます。AtCoder Problems による Difficulty は、1363 でした。

解答案

入力例1について計算する

入力例1の給料の期待値を計算してみましょう。サイコロは、前に出た目以下になるまで振ります。出た目が3の場合は、それ以上は振りません。次に出る目は、必ずそれ以下になるためです。

和の期待値は、期待値の和になります。給料の期待値は以下となります。（10/27 + 11/27 = 21/27 = 7/9 となります）

$$\frac{3}{9}+\frac{10}{27}+\frac{11}{27}+\frac{9}{9}+\frac{4}{9}+\frac{8}{9}+\frac{18}{9}=\frac{49}{9}$$

漸化式を考える

後ろから計算してみましょう。具体的には、dp(i) を i が出た以降の給料の期待値とします。まず 3 がでれば、給料は 6 で確定します。

dp(3) から dp(2) を求めてみます。2 が出た場合の期待値は2となります。またこれ以降に3がでる確率は、1/3となります。その結果、dp(2)は以下となります。

dp(2) と dp(3) から dp(1) を求めてみます。1 が出た場合の期待値は3となります。これ以降に2または3がでる確率はそれぞれ1/3となります。加える期待値は、4/3+6/3 =(4+6)/3となります。

dp(1)、dp(2)、dp(3) は、それぞれ1/3の確率で発生します。このため最終的な期待値は、以下となります。

$$\left(\frac{19}{3} + 4 + 6 \right) / 3 = \frac{49}{9}$$

上で考察した漸化式は以下となります。sum は、dp[j] (i < j) の和となります。dp は、Nから1に向かって値を決めていきます。

• dp[i] = a[i] + sum / N
• sum += dp[i]

初期値は、sum = 0 となります。最終的な期待値は、sum / N となります。

C++ プログラム例（ABC326E）

先に考察した漸化式を mod 998244353 で計算する必要があります（33ー36行目）。

MOD付の演算は、足し算、引き算、掛け算は自然にできますが、割り算は工夫が必要です。フェルマーの小定理を使えば、$N \pmod M$ で割ることは、$N^{M-2} \pmod M$ を掛けることと同じになります。

この計算をするために累乗を計算する関数 power を使います（7ー17行目）。この関数は自前ライブラリとして整備したものです。$N$ で割る場合に掛ける値を変数 invN として計算しておきます（28行目）。

以下が、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;
	cin >> n;
	vector<ull> a(n + 1, 0);
	for (int i = 1; i <= n; ++i) {
		cin >> a[i];
	}

	ull invN = power(n, MOD - 2, MOD);

	vector<ull> dp(n + 1, 0);
	ull sum = 0;
	for (int i = n; i >= 1; --i) {
		dp[i] = a[i] + sum * invN;
		dp[i] %= MOD;
		sum += dp[i];
		sum %= MOD;
	}

	cout << (sum * invN)% MOD << endl;

	return 0;
}

AtCoder では、AtCoder Library（ACL）と呼ばれるライブラリが提供されています。

ACL で、MOD付きの計算ができる型 modint も提供されています。この問題は、modint を使うとすっきりとプログラムを書くことができます（20、21行目）。以下は、modint を使ったプログラムです。

#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;

typedef modint998244353 mint;

int main()
{
	int n;
	cin >> n;
	vector<int> a(n + 1, 0);
	for (int i = 1; i <= n; ++i) {
		cin >> a[i];
	}

	vector<mint> dp(n + 1, 0);
	mint sum = 0;
	for (int i = n; i >= 1; --i) {
		dp[i] = mint(a[i]) + sum / mint(n);
		sum += dp[i];
	}

	sum /= mint(n);
	cout << sum.val() << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python プログラム例（ABC326E）

基本的な考え方は同じです。MOD計算付きの累乗は、組込み関数 pow で計算することができます（7行目）。

"""AtCoder Beginner Contest 326 E"""
MOD = 998244353
n = int(input())
a = list(map(int, input().split()))
a.insert(0, 0)

invN = pow(n, MOD - 2, MOD)

dp = [0] * (n + 1)
s = 0
for i in range(n, 0, -1):
    dp[i] = a[i] + s * invN
    dp[i] %= MOD
    s += dp[i]
    s %= MOD

print((s * invN) % MOD)

こちらも「AC」と判定されました。

最後に

確率を用いるDPは、ABC323E問題（解説記事）で初めて解くことができました。この問題は、解説動画で理解した内容を記事にしました。

引き続き ABC の問題を紹介していきます。