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

Правка en4, от aniervs, 2024-07-30 08:16:54

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:

$$$\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:

$$$\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$$$
$$$\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$$$
$$$(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:

$$$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:

$$$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, S[l...mid])
    update S[mid+1...r] with the convolution's results
    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 [0, 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 as well.

Теги convolutions, math, divide and conquer

История

 
 
 
 
Правки
 
 
  Rev. Язык Кто Когда Δ Комментарий
en10 Английский aniervs 2024-07-30 15:24:58 11 Tiny change: 'nvolve(A, S[l...mid])\n upd' -> 'nvolve(A, B)\n upd'
en9 Английский aniervs 2024-07-30 08:36:42 27
en8 Английский aniervs 2024-07-30 08:30:53 2 Tiny change: 's $j \in [0, K+1]$, h' -> 's $j \in [1, K+1]$, h'
en7 Английский aniervs 2024-07-30 08:19:25 81 Tiny change: ':\n\n\n$$\sum_{i=1}' -> ':\n\n\n$$\displaystyle\sum_{i=1}'
en6 Английский aniervs 2024-07-30 08:17:38 2 (published)
en5 Английский aniervs 2024-07-30 08:17:16 10 Tiny change: 'ough them as well.\n\n' -> 'ough them too.\n\n'
en4 Английский aniervs 2024-07-30 08:16:54 199
en3 Английский aniervs 2024-07-30 08:12:42 331
en2 Английский aniervs 2024-07-30 08:08:08 366 Tiny change: ' like this:\n\n~~~~~' -> ' like this, in a very Pythonic style:\n\n~~~~~'
en1 Английский aniervs 2024-07-30 07:58:28 2774 Initial revision (saved to drafts)