AtCoder が提供しているABC（AtCoder Beginner Contest）323 のE問題をC++とPythonで解いてみました。ABC323は、2023年10月7日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Playlist（Difficulty : 1279）
問題はリンク先をご覧ください。
確率DPで解くことができました。AtCoder Problems による Difficulty は、1279 でした。※2023年10月15日 Difficulty を追記しました。
解答案
確率を求める問題です。背反な事象のいずれかが発生する確率は、それぞれの事象の確率の和となります（和の法則）。
dp[i] を時刻 $i$ で曲が切り替わる確率とします。dp[i] は以下の式で求めることができます。以下の式で $i – t_j \;(0 \leqq j < n)$ が $0$ 以上の場合のみ加えます。確率の和になっているのは、それぞれが背反事象であるからです。
dp[i] = dp[i – t0] / N + … + dp[i – tN] / N
初期値は、
dp[0] = 1、それ以外は dp[i] = 0
とします。曲が切り替わる確率を求めることができれば、時刻 $X+0.5$ 秒後で、曲0 が流れている確率は、$0 \leqq i \leqq X$ における以下の和となります。
dp[i] / N (i + t0 > X の場合のみ加える)
確率の和になっているのは、それぞれの事象が背反であるからです。
C++ プログラム例（ABC323E）
MOD付の演算は、足し算、引き算、掛け算は自然にできますが、割り算は工夫が必要です。フェルマーの小定理を使えば、$N \pmod M$ で割ることは、$N^{M-2} \pmod M$ を掛けることと同じになります。
この計算をするために累乗を計算する関数 power を使います（7ー17行目）。この関数は自前ライブラリとして整備したものです。$N$ で割る場合に掛ける値を変数 invN として計算しておきます（28行目）。
先に考察した漸化式を計算します（35、44行目）。
以下が、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, x;
cin >> n >> x;
vector<int> t(n);
for (int i = 0; i < n; ++i) {
cin >> t[i];
}
ull invN = power(n, MOD - 2, MOD);
vector<ull> dp(x + 1, 0);
dp[0] = 1;
for (int i = 1; i <= x; ++i) {
for (int j = 0; j < n; ++j) {
if (i - t[j] >= 0) {
dp[i] += (dp[i - t[j]] * invN) % MOD;
dp[i] %= MOD;
}
}
}
ull result = 0;
for (int i = 0; i <= x; ++i) {
if (i + t[0] > x) {
result += (dp[i] * invN) % MOD;
result %= MOD;
}
}
cout << result << endl;
return 0;
}AtCoder では、AtCoder Library（ACL）と呼ばれるライブラリが提供されています。
ACL で、MOD付きの計算ができる型 modint も提供されています。この問題は、modint を使うと簡潔にプログラムを書くことができます。
- 必要な宣言は、2行目、4行目、6行目です。
- modint を使うことにより漸化式がすっきり書けています（22、30行目）。
以下は、modint を使ったプログラムです。
#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
typedef modint998244353 mint;
int main()
{
int n, x;
cin >> n >> x;
vector<int> t(n);
for (int i = 0; i < n; ++i) {
cin >> t[i];
}
vector<mint> dp(x + 1, 0);
dp[0] = 1;
for (int i = 1; i <= x; ++i) {
for (int j = 0; j < n; ++j) {
if (i - t[j] >= 0) {
dp[i] += dp[i - t[j]] / n;
}
}
}
mint result = 0;
for (int i = 0; i <= x; ++i) {
if (i + t[0] > x) {
result += dp[i] / n;
}
}
cout << result.val() << endl;
return 0;
}どちらも AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC323E）
基本的な考え方は同じです。累乗は演算子 ** で計算できます。MOD計算付きの累乗は、組込み関数 pow で計算することができます（6行目）。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 323 E"""
MOD = 998244353
n, x = map(int, input().split())
t = list(map(int, input().split()))
invN = pow(n, MOD - 2, MOD)
dp = [0] * (x + 1)
dp[0] = 1
for i in range(1, x + 1):
for j in range(n):
if i - t[j] >= 0:
dp[i] += (dp[i - t[j]] * invN) % MOD
dp[i] %= MOD
result = 0
for i in range(x + 1):
if i + t[0] > x:
result += (dp[i] * invN) % MOD
result %= MOD
print(result)最後に
確率を用いるDPは苦手意識があって、解く前にあきらめていました。公式解説を読んで、漸化式を理解することができました。
確率DPは、一定の頻度で出題されるため、過去問を解こうと思います。
引き続き ABC の問題を紹介していきます。
