2023-05-15

AtCoder ARC 160 D - Mahjong (橙色, 700 点)

AtCoder AtCoder700点 ARC-D 多項式・FPS(形式的冪級数) Bostan-Mori法制約条件:総和=K 二項係数負の二項係数入力が定数個対象を一意に定める操作列を数え上げる数え上げ問題グラフ・盤面・数列の個数の数え上げ操作後の結果の数え上げ操作橙色diff FPS(形式的冪級数)の高等演算 FPSテク:漸化式の母関数を求める重複組合せ数列 Greedy：先に進むほど新たな選択肢が挿入される包除原理制約条件:xi<=ai

FPS による考察はやっぱり強力ですね！

問題へのリンク

問題概要

長さ $N$ かつ総和 $M$ である非負整数列 $A=(A_{1}, A_{2}, \dots, A_{N})$ のうち、以下の条件を満たすものの個数を 998244353 で割ったあまりを求めよ。

「以下の操作のうちどちらかを選んで行うことを繰り返して、 $A$ の全ての要素を 0 にすることが出来る。」

操作 1： $1 \le i \le N$ を満たす整数 $i$ を選び、 $A_{i}$ を $K$ 減らす
操作 2： $1 \le i \le N−K+1$ を満たす整数 $i$ を選び、 $A_{i}, A_{i+1}, \dots, A_{i+K-1}$ を 1 ずつ減らす

制約

$1 \le K \le N \le 2000$
$1 \le M \le 10^{18}$

考えたこと：操作列の数え上げに帰着する

操作によって出来上がるものが何通りあるかを数え上げる問題では、大きく分けて

操作によって作れるかどうかの判定問題をまず解く
各対象物を作るための操作列が一意となるようにして、その操作列を数え上げる

という方針があるように思える。今回は後者の考え方が適用できる。

各 $i$ に対して、操作 2 の適用回数は $K$ 回未満と思った良いことに注意する。操作 2 を $i$ に対して $K$ 回適用することは、操作 1 を $i, i+1, \dots, i+K-1$ に対して 1 回ずつ適用することで代用できるからだ。

そして、どの $i$ に対しても操作 2 の適用回数が $K$ 回未満であるように限定した場合に得られる数列に対しては、必要な操作が一意に決まることが言える。具体的には、「左端の 0 より大きいところから Greedy に消していく」という方針によって、操作列が一意に決まる。

さらに整理すると、次のことが言える。

操作列において、操作の順序は入れ替えても同一視する
各 $i$ に対して、操作 2 は $K$ 回未満とする
各 $i$ に対して、操作 1 は任意回数実行する

として得られる数列を考えると、異なる操作列からは異なる数列が生まれることとなる。

よって、操作列を数え上げればよいことになった。

操作列の数え上げ

大前提として $M$ が $K$ の倍数でない場合は 0 通りとなる。よって、 $M$ は $K$ の倍数とする。このとき、操作は $M' = \displaystyle \frac{M}{K}$ 回行うことになる。

整理すると、私たちは

$0 \le s_{1}, s_{2}, \dots, s_{N-K+1} \le K-1$
$t_{1}, t_{2}, \dots, t_{N} \ge 0$
$s_{1} + s_{2} + \dots + s_{N-K+1} + t_{1} + t_{2} + \dots + t_{N} = M'$

を満たすような整数 $s_{1}, s_{2}, \dots, s_{N-K+1}, t_{1}, t_{2}, \dots, t_{N}$ の組の個数を求めたいことになる。ここから先はさまざまな解法が考えられる。

解法 1：FPS で考えて、Bostan-Mori 法を適用する

一般に「総和が $M'$ 」などといった制約を見ると FPS が適用できそうに思える。実際、今回も

変数 $s_{i}$ に対応する FPS： $1 + x + x^{2} + \dots + x^{K-1} = \displaystyle \frac{1 - x^{K}}{1 - x}$
変数 $t_{i}$ に対応する FPS： $1 + x + x^{2} + \dots = \displaystyle \frac{1}{1 - x}$

これらの積をとって、 $M'$ 次の項を求めればよい。

$\displaystyle \bigl(\frac{1 - x^{K}}{1 - x} \bigr)^{N-K+1} \bigl(\frac{1}{1 - x}\bigr)^{N} = \frac{(1 - x^{K})^{N-K+1}}{(1 - x)^{2N-K+1}}$

であるから、最終的に次の値が求められればよいことがわかった。

$\displaystyle \lbrack x^{M'} \rbrack \frac{(1 - x^{K})^{N-K+1}}{(1 - x)^{2N-K+1}}$

これは、Bostan-Mori 法 (この記事に詳しく書いた) を適用することで、 $O(KN \log N \log M)$ の計算量で求められる。

提出コード (Bostan-Mori 法)

分子の $(1 - x^{K})^{N-K+1}$ と分母の $(1 - x)^{2N-K+1}$ については、二項定理を用いて展開した式を直接入力した。

atcoder.jp

解法 2：負の二項定理を使う

Bostan-Mori 法をそのまま適用するのは、少し計算時間に不安がある。そこで、もう少し考察を進めることにする。

$\displaystyle \lbrack x^{M'} \rbrack \frac{(1 - x^{K})^{N-K+1}}{(1 - x)^{2N-K+1}}$

という式の分母分子を二項定理で展開すると、次のようになる。なお、分母については $(1 - x)^{-(2N-K+1)}$ と考えて、機械的に二項定理を適用する。

$\displaystyle \lbrack x^{M'} \rbrack \frac{(1 - x^{K})^{N-K+1}}{(1 - x)^{2N-K+1}}$
$= \displaystyle \sum_{i = 0}^{N-K+1} (-1)^{i} C(N-K+1, i) \lbrack x^{M' - iK} \rbrack \frac{1}{(1 - x)^{2N-K+1}}$
$= \displaystyle \sum_{i = 0}^{N-K+1} (-1)^{i} C(N-K+1, i) (-1)^{M'-iK} C(-(2N-K+1), M'-iK)$