解きます
解く
分割統治典型です. $k$次の項の係数が「$k$個の数を選んだときの総和」を表すことにすると, 求めるものは
$$\left[x^K\right]\prod_{k=1}^N\left(1+A_kx\right)$$
になります. 愚直にやると$O\left(N^2\right)$ですが, これを$O\left(N\log^2N\right)$で求められます.
前提として, 多項式積はFFTを用いることで$O\left(N\log N\right)$で計算できます. ここで$N$は「求めたい最大次数」です. この方法については調べたら沢山記事が出てくるし, 去年私も記事を書いたので解説しません.
この手法が活きるのは「掛ける2つの多項式の最高次数が近いとき」です. 愚直にやるだけでは計算をすすめるにつれて「(ある程度大きい多項式)$\times$(1次式)」みたいな形になってしまい, FFTを活かせません.
これを改善するために, 以下のようなことをします.
- 数列$A=(A_1,A_2,...,A_N)$を$B=(A_1,A_2,...,A_{N\div2})$と$C=(A_{N\div2+1},A_{N\div2+2}m,...,A_N)$の2つの数列に分割する
- $B$, $C$でそれぞれ積を求める
- 求めた2つをFFTで畳み込む
これを再帰的に行うことで, $O\left(N\log^2 N\right)$になります.
計算量の計算
簡単のため, $N$が2の羃と仮定します.
行う積を列挙すると以下のようになります
- $N/2$次同士の積が1回
- $N/4$次同士の積が2回
- $N/8$次同士の積が4回
- ...
- $2$次同士の積が$N/4$回
- $1$次同士の積が$N/2$回
$N=2^B$とすると, 計算量は
$$\begin{align}
O\left(\sum_{k=0}^{B-1}2^{B-1-k}\cdot 2^k\log 2^k\right)
&=O\left(2^{B-1}\sum_{k=0}^{B-1}\log 2^k\right) \\
&=O\left(2^{B-1}\left(B-1\right)\log 2^{B-1}\right) \\
&=O\left(N\log^2N\right) \\
\end{align}$$
となります. ($L=O\left(\log N\right)$です)
use ac_library::*;
use proconio::*;
type Mint = ModInt998244353;
// (1+a[0]*x)*(1+a[1]*x)*(1+a[2]*x)*... を計算する
fn solve(a: &[Mint]) -> Vec<Mint> {
// 空積は1
if a.is_empty() {
return vec![Mint::raw(1)];
}
// 1つの要素しかないなら, 1+ax
if a.len() == 1 {
return vec![Mint::raw(1), a[0]];
}
// 数列の真ん中
let mid = a.len() / 2;
// 左半分の結果
let l = solve(&a[..mid]);
// 右半分の結果
let r = solve(&a[mid..]);
// AC Libraryの畳み込みを使用
convolution(&l, &r)
}
fn main() {
input! {
n: usize,
k: usize,
a: [Mint; n],
}
println!("{}", solve(&a).get(k).unwrap_or(&Mint::raw(0)));
}
提出(170 ms):
関連問題: これを多倍長整数を使って実際に値を求めることで解いてみましょう
終わり
今日は典型問題って感じがしましたね
以上です