解きます
解く
サイコロに対応する関数を
$$d := \frac 1M\sum_{k=1}^M x^{A_M}$$
とします.
まず, ゴールも落とし穴が無い場合, $k$次の係数が「マス$k$を通る確率」になるfpsは
$$\sum_{k=0}^\infty d^n = \frac1{1-d}$$
になります. 例えばここでマス$B_0$にだけ落とし穴があるとすると, $B_0$を通る組合せだけ消してやればいいので,
$$\frac1{1-d}-\left(\left[x^{B_0}\right]\frac1{1-d}\right)\frac{x^{B_0}}{1-d}$$
となります. 複雑に見えますが, $\left[x^{B_0}\right]f=0$になるように$\frac1{1-d}-X\frac{x^{B_0}}{1-d}$の定数$X$を決定した, というだけです.
このノリで, $B_0,...,B_L$に穴がある場合は,
$$f=\frac1{1-d}-X_0\frac{x^{B_0}}{1-d}-X_1\frac{x^{B_1}}{1-d}-\cdots-X_L\frac{x^{B_L}}{1-d}$$
のようになり, $\left[x^{B_0}\right]f=\left[x^{B_1}\right]f=\cdots=\left[x^{B_L}\right]f=0$になるように係数$X_0, X_1,\dots X_L$を決定すればいいです.
「ゴールを行きすぎてもゴールで止まる」はFPSに落とし込むのが大変そうなので, マス$N-1$までの確率を求めた後, マス$N$だけ愚直に求めるようにしましょう.
オンライン畳み込み
ここでRelaxed Convolution(オンライン畳み込み)というものを使います.
以下のようなことが$O\left(N\log^2N\right)$みたいな計算量でできます. (実際はもっと条件を緩めたりできますが, 今回使う分はこれだけです)
- 不明な多項式$A$と分かってる多項式$B$がある
- 以下を上から順番に行う
- $\left[x^0\right]A$を入力し, $\left[x^0\right]A B$を出力する
- $\left[x^1\right]A$を入力し, $\left[x^1\right]A B$を出力する
- $\left[x^2\right]A$を入力し, $\left[x^2\right]A B$を出力する
- $\left[x^3\right]A$を入力し, $\left[x^3\right]A B$を出力する
- ...
この記事の図を見たら仕組みは一瞬で理解できます. 今回使うのはSemi-relaxed Convolutionですね.
今回は$\left[x^n\right]A$を入力する前に「$\left[x^n\right]A=0$だった場合の$\left[x^n\right]AB$」を知る必要がありますが, 1次先読みするくらいならどうとでもなります. (極論, $\left[x^n\right]A=0$で実際に計算したあとそれを打ち消せばいいだけですし)
ニュートン法
実装するには$\frac1{1-d}$の計算の仕方も知る必要があります. これにはニュートン法というものを使います.
ニュートン法は汎用的な方程式の解を求める手法なのですが, これで$1/h$を求めるには,
$z_{n+1} = 2z_n-hz_n^2$
という漸化式を反復すればいいです.
この収束の速さを見てみます. $z_n=1/h+\epsilon$とすると
$\begin{align}
z_{n+1}
&= 2\left(\frac1h+\epsilon\right)-h\left(\frac1h+\epsilon\right)^2 \\
&= \frac2h+2\epsilon-h\left(\frac1{h^2}+2\frac\epsilon{h}+\epsilon^2\right) \\
&= \frac1h-h\epsilon^2 \\
\end{align}$
と, おおよそ誤差が2乗になります.
これ, つまり逆数が$k-1$次の項まで求まっている(誤差が$k$次以上)のとき, この反復1回で$2k-1$次まで求まる(誤差が$2k$次以上)ということです.
一回の反復で次数が2倍になるので, $O(\log N)$回の反復で済み, 計算量は$O(N\log N)$になります.
ちょっと待って!
漸化式は掛け算を含んでいて, FFT畳み込みを使って$O\left(N\log N\right)$かかります.
それを$O\left(\log N\right)$回繰り返すので$O\left(N\log^2N\right)$かかると考えた人がいるかもしれません.
計算する際は「精度を倍々にしていく」のでそれを考慮して計算量を見ると
$$\begin{align}
O\left(\sum_{k=0}^{\log_2 N}\frac N{2^k}\log\frac N{2^k}\right)
&= O\left(\sum_{k=0}^{\log_2 N}\frac N{2^k}\log N\right) \\
&= O\left(2N\log N\right) \\
&= O\left(N\log N\right) \\
\end{align}$$
となります. 反復の最初の方はそんなに精度を求めないから計算量として無視できる程度に小さい, ということですね.
実装
必要な知識は全て揃いました. あとはコードを書くだけです. オンライン畳み込みの方が支配的で, 全体で$O\left(N\log^2N\right)$です.
use ac_library::*;
use proconio::*;
type Mint = ModInt998244353;
fn main() {
input! {
n: usize,
m: usize,
l: usize,
a: [usize; m],
b: [usize; l],
}
// mの逆数
let minv = Mint::from(m).inv();
// 1-dの逆数 (n-1次まで)
let z = {
// h = 1-dとする. 長さはニュートン法の都合で2の羃に伸ばす
let mut h = vec![Mint::raw(0); n.next_power_of_two()];
h[0] = Mint::raw(1);
for &a in a.iter() {
// n-1次まで求めるので, n次は不要
if a == n {
continue;
}
h[a] = -minv;
}
// 1で初期化. この時点で誤差は1次
let mut z = vec![Mint::raw(1)];
// 十分な精度になるまで反復する
while z.len() < h.len() {
// FFTの性質を上手く使えばもっと最適化できますが, このままでも十分速いです
let m = convolution(&z, &z);
let m = convolution(&m, &h[..z.len() * 2]);
z.extend((0..z.len()).map(|_| Mint::raw(0)));
for i in 0..z.len() {
z[i] = z[i] * 2 - m[i];
}
}
z
};
// 穴の補正分の多項式x. 定数関数1で初期化
let mut x = vec![Mint::raw(1)];
// 目的の多項式f
let mut f = vec![Mint::raw(0); n];
f[0] = Mint::raw(1);
// 配列bを読む進捗管理用
let mut k = 0;
// オンライン畳み込みの実行
for i in 1..n {
// xの係数の決定
x.push(if b.get(k) == Some(&i) {
// 穴がある場合
k += 1;
// 穴が無かった場合のf[i]を計算
let mut v = f[i];
// まだ畳み込んで無い部分を付ける
for j in 0..=i.trailing_zeros() {
v += x[i - (1 << j)] * z[1 << j];
}
// それを打ち消すように係数を決定
-v
} else {
// 穴が無いので補正も無い
Mint::raw(0)
});
// 畳み込みを計算
f[i] += x[i] * z[0];
for j in 0..=i.trailing_zeros() {
let h = convolution(&x[i - (1 << j)..i], &z[1 << j..2 << j]);
for (l, h) in h.into_iter().enumerate() {
// 必要な分を越えたら終了
if l + i >= n {
break;
}
f[l + i] += h;
}
}
}
let mut ans = Mint::raw(0);
let mut k = 0;
for i in (0..n).rev() {
// 「マスiからゴールするにはA_k以上の出目が必要」になるようkを調整
while k < m && i + a[k] < n {
k += 1;
}
// そのkに応じて確率を足す
ans += minv * f[i] * (m - k);
}
println!("{ans}");
}
提出(413 ms):
終わり
盛りだくさんでしたね.
以上です