Contest Duration: - (local time) (120 minutes) Back to Home
Official

## 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.

posted:
last update: