An unexpected D&C with FFT

Revision en1, by Um_nik, 2022-06-07 22:37:28

I came up with this idea when solving an AtCoder problem, but the editorial had a simpler solution. I liked the idea, so I decided to set a CodeChef problem using this idea. Since most of you guys don't participate in CodeChef contests, which is a big mistake, I'll write a short blog explaining the idea. Consider this a spoiler warning for these two problems and a CodeChef promotion.

Combinatorial problem

For given $a$ and $n \le 10^5$ calculate $A_k = \sum_{i=0}^{k} \binom{k}{i} \binom{n-k}{k-i} a^{i}$ for each $k$ from $0$ to $n$.

(Use the two links above to see why anybody would like to do that)

Reformulation with polynomials

It is easy to see that $A_k$ is $((1+ax)^{k}(1+x)^{n-k})[x^k]$ (coefficient of $x^k$ in polynomial $(1+ax)^{k}(1+x)^{n-k}$). But since the polynomials are different for different $k$ this doesn't seem very helpful...

Solution

Let's generalize this a bit and say that $P$ and $Q$ are polynomials of degree at most $d$ and we want to calculate $P^k Q^{n-k} [x^k]$ for all $k$ ($d = 1$ in the problem above).

Let's apply divide-and-conquer. Let's say we want to solve the problem for $k \in [l, r]$. Then all the interesting polynomials will be divisible by $P^l Q^{n-r}$. Let's say we somehow calculated this polynomial. For given $k$ we will then have to multiply it by $P^{k-l} Q^{r-k}$. But the degree of this polynomial is $d(r-l)$, and since in the end we will be looking only at coefficients from segment $[l, r]$, right now we only care about the coefficients from segment $[l - d(r-l), r]$, all the other coefficients can't affect the answer. So, the number of coefficients we are interested in is $O(d(r-l))$. Let's write a recursive function that will take $l$, $r$ and a segment of interesting coefficients of $P^l Q^{n-r}$ and will find all the answers for $k \in [l, r]$. We'll choose $m = \frac{l+r}{2}$ and recurse to $[l, m]$ and $[m + 1, r]$. To recurse to the left, for example, we'll need to multiply $P^l Q^{n-r}$ by $Q^{r-m}$ and take a segment of coefficients. Since both polynomials have degree $O(d(r-l))$, we can multiply them in $O(d(r-l) \log (d(r-l)))$ time. To get $Q^{r-m}$ one can use binary exponentiation, it will work in $O(d(r-m) \log (d(r-m)))$ time (if $d=1$, one can use binomial theorem instead). The total complexity will be $O(nd \log^{2} n)$.

History

Revisions

Rev. Lang. By When Δ Comment
en1 Um_nik 2022-06-07 22:37:28 2486 Initial revision (published)