F - Modulo Sum of Increasing Sequences Editorial by evima

For convenience, we denote \(\mathrm{MOD}\) by \(P\).

[1] Rephrasing the problem

Consider the sequence \(B\) defined by \(B_i \coloneqq A_i + i\), and one can see that it corresponds one-to-one with \(A\), so let us count \(B\) instead of \(A\).

Here, \(B\) is strictly increasing, and its elements take values between \(0\) and \(M+N-1\), so \(B\) corresponds one-to-one with a size-\(N\) subset of \(\{0,1,\ldots,M+N-1\}\).

From these arguments, if we redefine \(M\) to be \(M+N\), the problem can be rephrased into the following:

How many ways are there to choose exactly \(N\) from \(0,1,\ldots,M-1\) so that \((\text{the sum}) \bmod P=K\), for each \(K=0,1,2,\ldots,P-1\)?

Below, we consider this version of the problem.

[2] Focusing on a special case

It is difficult to solve the problem above for all cases directly, so let us first consider the case \(M \bmod P=0\).

Let \(V_s\) be the number of choices with the sum of \(s\) in the problem above, and consider the polynomial \(f(x) = \sum_s V_s x^s=[y^N]\prod_{i=0}^{M-1} (1+x^iy)\).

Finding \(g(x) = f(x) \bmod (x^P-1)\) will solve the problem.

[3] Using the principles of Fourier transform

Let \(\omega\) be the primitive \(P\)-th root of \(1\) in \(\mathbb{C}\) (the complex field). If \(g(\omega^0),g(\omega^1),\ldots,g(\omega^{P-1})\) are found, \(g\) can be restored in a way similar to inverse Fourier transform. (We have complex numbers here, but eventually, we will only have to work with rational numbers.)

Here, we have \(g(\omega^k) = f(\omega^k)\), so let us find \( f(\omega^k)\) instead. A substitution yields \(f(\omega ^k)=[y^N]\prod_{i=0}^{M'-1}(1+(w^k)^i y)\).

Now, let us use the assumption that \(M \bmod P=0\).

Define \(H(y)=\prod_{i=0}^{P-1}(1+(w^k)^i y)\), and we have \(f(\omega ^k)=[y^N] H(y)^{M/P}\).

Here, let us rewrite \(H\). Let \(g\coloneqq \mathrm{gcd}(k,P)\). Using the factor theorem, one computes that \(H(y)=(1-(-y)^{\frac{P}{g}})^g\).

Therefore, \(f(\omega ^k)=[y^N] (1-(-y)^{\frac{P}{g}})^{\frac{gM}{P}}\), which can be computed using the binomial theorem.

Now that \(g(\omega^0),g(\omega^1),\ldots,g(\omega^{P-1})\) are found, let us restore \(g(x)\).

It holds that \([x^j]g(x)=\frac{1}{P}\sum_{i=0}^{P-1}\frac{g(\omega^i)}{\omega^{ij}}\).

Directly applying inverse Fourier transform to the above formula requires division by the complex number \(\omega\). To avoid this, let us use the fact that \(g(\omega^a)=g(\omega^b)\) if \(\mathrm{gcd}(a,P)=\mathrm{gcd}(b,P)\) to put together the computations where the values of \(g\) are equal.

One can compute \(\sum_{0\le i < P,\mathrm{gcd(i,P)=k}}\frac{1}{\omega^{ij}}\) using the inclusion-exclusion principle, for example, which also shows that this value is an integer. Additionally, \(g(\omega^0),g(\omega^1),\ldots,g(\omega^{P-1})\) are all rational numbers, so all computations can be done with rational numbers without complex numbers.

Therefore, the special case can be solved in \(\mathrm{O}(P^2)\) time after pre-computing factorials in \(\mathrm{O}(N+M)\) time.

[4] Solving the general case

Consider the case \(M \bmod P\neq 0\). Let \(M'\) be the greatest multiple of \(P\) not exceeding \(M\), and fix the number of chosen values between \(M'\) and \(M-1\).

The choices of values between \(0\) and \(M'-1\) can be counted in \(\mathrm{O}(P^2)\) time using the method in [3].

For the choices of values between \(M'\) and \(M-1\), the problem is to find the number of ways to choose exactly \(i\) values between \(M'\) and \(M-1\) so that \((\text{the sum}) \bmod P=j\), which can be solved with simple dynamic programming.

Then, one can straightforwardly merge these two parts.

The dynamic programming part needs to be pre-computed just once, which, combined with the fact \(M-M'<P\), enables one to try all possible numbers of chosen values between \(M'\) and \(M-1\) in \(\mathrm{O}(P^3)\) time.

Therefore, the problem is solved in \(\mathrm{O}(N+M+P^3)\) time.

last update: