Convolution present on 1^k + 2^k + ... + n^k

#	User	Rating
1	tourist	4009
2	jiangly	3823
3	Benq	3738
4	Radewoosh	3633
5	jqdai0815	3620
6	orzdevinwang	3529
7	ecnerwala	3446
8	Um_nik	3396
9	ksun48	3390
10	gamegame	3386

#	User	Contrib.
1	cry	167
2	Um_nik	163
3	maomao90	162
3	atcoder_official	162
5	adamant	159
6	-is-this-fft-	158
7	awoo	157
8	TheScrasse	154
9	Dominater069	153
9	nor	153

The following is a very well-known problem, yet I recently found a nice solution involving convolutions.

You're given $$$K$$$ and $$$n$$$ and you must compute for all $$$k \in [0, K]$$$

$$$S_k = 1^k + 2^k + \dots + n^k$$$

Obviously, we want it $$$mod$$$ some $$$M$$$, but let's ignore it for now.

There're many ways of solving this problem, but here I focus on a particular one:

We start with some standard stuff, by expanding $$$\sum_{i=1}^n (i+1)^{k + 1}$$$. Through the Binomial Theorem we get:

$$$\displaystyle\sum_{i=1}^n (i+1)^{k + 1} = \sum_{i=1}^n \sum_{j=0}^{k+1} \binom{k+1}{j} i^j = \sum_{j=0}^{k+1} \binom{k+1}{j} \sum_{i=1}^n i^j = \sum_{i=1}^n i^{k + 1} + \binom{k + 1}{k} S_{k}+ \sum_{j=0}^{k-1} \binom{k+1}{j} S_j$$$

Then, we can actually cancel out the LHS via some telescopic sum:

$$$\displaystyle\sum_{i=1}^n (i+1)^{k + 1} = \sum_{i=1}^n i^{k + 1} + \binom{k + 1}{k} S_{k}+ \sum_{j=0}^{k-1} \binom{k+1}{j} S_j$$$ $$$\displaystyle\sum_{i=1}^n (i+1)^{k + 1} - \sum_{i=1}^n i^{k + 1} = \binom{k + 1}{k} S_{k}+ \sum_{j=0}^{k-1} \binom{k+1}{j} S_j$$$ $$$\displaystyle (n+1)^{k + 1} - 1 = (k + 1) S_{k}+ \sum_{j=0}^{k-1} \binom{k+1}{j} S_j$$$

From there, we get a recurrence:

$$$\displaystyle S_{k} = \frac{1}{k + 1} \left[(n+1)^{k + 1} - 1 - \sum_{j=0}^{k-1} \binom{k+1}{j} S_j \right]$$$

Now, let's expand the binomial coefficients and we get:

$$$\displaystyle S_{k} = \frac{1}{k + 1} \left[(n+1)^{k + 1} - 1 - (k+1)! \cdot \sum_{j=0}^{k-1} \frac{1}{(k+1-j)!} \cdot \frac{1}{j!} S_j \right]$$$

Now, let's define for each $$$t$$$: $$$A_{t}$$$ as $$$\frac{1}{(t+1)!}$$$ and $$$B_{t}$$$ as $$$\frac{1}{t!} S_t$$$. Notice that now the inner sum in the recurrence it's almost a convolution:

We have pairs $$$A_{i}, B_{j}$$$, and the sum of the products $$$A_i \cdot B_j$$$ is contributing to $$$S_{i+j}$$$. The main issue is that we can't do a single convolution between $$$A$$$ and $$$B$$$ in order to compute $$$S$$$, because we need $$$S$$$ to define $$$B$$$.

The alternative is to do some Divide and Conquer (thanks to Errichto who taught me this trick 2 years ago or more):

Let's say we want to compute $$$S_k$$$ for every $$$k \in [l, r]$$$. Let's say $$$p=\lfloor \frac{l+r}{2} \rfloor$$$. We'll have a recursive algorithm:

The idea is to first compute $$$S_k$$$ for each $$$k \in [l, p]$$$ recursively. Then to compute $$$\mathrm{B}[]$$$ for those values of $$$k$$$ and to apply the convolution. Then to update $$$S_k$$$ for all $$$k \in [p+1, r]$$$. Lastly, solve recursively for the other half ($$$[p+1,r]$$$).

Something like this, in a very Pythonic style:

def solve(l, r):
    if l == r:
        # base case, add the parts not related to the convolution
        # use proper modular arithmetics 
        # ...
        return
    mid = (l + r) // 2
    solve(l, mid)
    convolve(A[1...r-l], B[l...mid])
    update S[mid+1...r] with the results from the convolution
    solve(mid + 1, r)

The final answer is computed by calling solve(0, K). Given that the sequences being convolved are of size $$$\mathcal{O}(r - l)$$$, we can convolve them in $$$\mathcal{O}((r - l)\log (r - l))$$$ time via FFT, giving a total time complexity of $$$\mathcal{O}(K\log^2 K)$$$.

As for modular arithmetics, notice we need the multiplicative inverse of all numbers $$$j \in [1, K+1]$$$, hence the modulo better be good (a big prime and also FFT-friendly).

That's all, I hope you like it. There are other ways of computing $$$S$$$, some of them quite interesting as well. I recommend going through them too.