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)$$$.