Official

F - Modulo Sum of Increasing Sequences Editorial by nok0


以下の解説では読みやすさのため、問題文中の \(\mathrm{MOD}\) のことを \(P\) と表記します。


[1] 問題の言い換え

\(B_i \coloneqq A_i + i\) で定義される数列 \(B\) を考えると、\(A\)\(B\) は一対一対応していることがわかります。よって、以下では \(A\) の代わりに \(B\) を数え上げることとします。

ここで、 \(B\) は狭義単調増加であり、\(0\) 以上 \(M+N\) 未満の値を取ります。このことから、 \(B\) は集合 \(\{0,1,\ldots,M+N-1\}\) のサイズ \(N\) の部分集合と一対一対応しています。

以上の議論より、 \(M\)\(M+N\) で置き直すと、問題は以下のように言い換えられます。


\(0,1,\ldots,M-1\) からちょうど \(N\) 個を選ぶとき、選んだ \(N\) 個の整数の総和 \(\bmod P=K\) となるのは何通りか、各 \(K=0,1,2,\ldots,P-1\) について求めよ。


以下では言い換え後の問題について考えます。


[2] 特殊な場合について解く

言い換え後の問題を直接全ての場合について解くのは難しいので、まずは \(M \bmod P=0\) の場合を考えましょう。

上の問題で総和が \(s\) になる場合の数を \(V_s\) として、多項式 \(f(x) = \sum_s V_s x^s=[y^N]\prod_{i=0}^{M-1} (1+x^iy)\) を考えます。

\(g(x) = f(x) \bmod (x^P-1)\) が求まればこの問題に答えられます。


[3] フーリエ変換の考え方の利用

\(\omega\)\(1\)\(\mathbb{C}\)(複素数体) における原始 \(P\) 乗根として、\(g(\omega^0),g(\omega^1),\ldots,g(\omega^{P-1})\) がわかれば逆フーリエ変換の要領で \(g\) が復元できます。(ここで複素数が出てきますが、実は全て有理数で処理できます)

ここで \(g(\omega^k) = f(\omega^k)\) が成立するので、以下 \( f(\omega^k)\) を求めます。代入して \(f(\omega ^k)=[y^N]\prod_{i=0}^{M-1}(1+(w^k)^i y)\) を得ます。

ここで、\(M \bmod P=0\) であることを使いましょう。

\(H(y)=\prod_{i=0}^{P-1}(1+(w^k)^i y)\) と置くと、\(f(\omega ^k)=[y^N] H(y)^{M/P}\) です。

ここで \(H\) を整理しましょう。\(g\coloneqq \mathrm{gcd}(k,P)\) と置きます。因数定理を用いると、\(H(y)=(1-(-y)^{\frac{P}{g}})^g\) と計算できます。

以上より \(f(\omega ^k)=[y^N] (1-(-y)^{\frac{P}{g}})^{\frac{gM}{P}}\) です。これは二項定理により求められます。

\(g(\omega^0),g(\omega^1),\ldots,g(\omega^{P-1})\) が求まったので、 \(g(x)\) を復元しましょう。

\([x^j]g(x)=\frac{1}{P}\sum_{i=0}^{P-1}\frac{g(\omega^i)}{\omega^{ij}}\) という式が成り立ちます。

上記の式でフーリエ逆変換を愚直に行うと複素数 \(\omega\) での除算が必要です。ここで \(\mathrm{gcd}(a,P)=\mathrm{gcd}(b,P)\) ならば \(g(\omega^a)=g(\omega^b)\) であることを用い、 \(g\) の値が等しいところをまとめて計算することとします。

\(\sum_{0\le i < P,\mathrm{gcd(i,P)=k}}\frac{1}{\omega^{ij}}\) は例えば約数包除で求められ、このことから整数であることも言えます。また \(g(\omega^0),g(\omega^1),\ldots,g(\omega^{P-1})\) は全て有理数でしたから、結局複素数を介さずに全て有理数上で計算できることがわかります。

よってこの問題は \(\mathrm{O}(N+M)\) で階乗前計算をしておくことで、 \(\mathrm{O}(P^2)\) で解けます。


[4] 一般の場合について解く

\(M \bmod P\neq 0\) の場合を考えます。\(M'\)\(M\) 以下の最大の \(P\) の倍数として、 \(M'\) 以上 \(M\) 未満の値を選ぶ個数を決め打ちます。

\(0\) 以上 \(M'\) 未満の部分は [3] の解法により \(\mathrm{O}(P^2)\) で解けます。

また、\(M'\) 以上 \(M\) 未満の部分についても、\(M'\) 以上 \(M\) 未満の値からちょうど \(i\) 個選ぶとき、総和 \(\bmod P\)\(j\) になるのは何通りかがわかればよいです。これは単純な動的計画法で求められます。

\(0\) 以上 \(M'\) 未満の部分と \(M'\) 以上 \(M\) 未満の部分のマージは愚直に行えば良いです。

動的計画法の部分は一回前計算しておけば十分なことから、\(M-M'<P\) なことと合わせて、結局全ての場合を決め打ちして試しても計算量は \(\mathrm{O}(P^3)\) です。

以上より、この問題を \(\mathrm{O}(N+M+P^3)\) で解くことができました。

posted:
last update: