COUNTIT - EDITORIAL

editorial
, ,
3

Here’s another solution that I think should work but I gave up trying to implement it since I figured it out around 1:00 AM the night before and I had never implemented FFT.

After you get the relation \sum_{x = 1}^{k} (x^m - (x - 1)^m) (x^n - (x - 1)^n) rewrite it as such

\sum_{x = 1}^{k} (x^m - (x - 1)^m) (x^n - (x - 1)^n) = \sum_{x = 0}^{k - 1} ((x + 1)^m - x^m) ((x + 1)^n - x^n)\\ = \sum_{x = 0}^{k - 1} (x + 1)^{m + n} + x^{m + n} - x^m(x + 1)^n - (x+1)^mx^n\\ = 2 S_{m + n} + k^{m + n} - \sum_{x = 0}^{k - 1} \left(x^m(x + 1)^n + (x+1)^mx^n \right)\\

where S_{r} = 1^r + 2^r + \dots + (k - 1)^r.
Now, looking at x^m(x + 1)^n

\sum_{x = 0}^{k - 1} x^m(x + 1)^n = \sum_{x = 0}^{k - 1} x^m\sum_{j = 0}^n \binom{n}{j} x^j = \sum_{j = 0}^n \binom{n}{j} \sum_{x = 0}^{k - 1} x^{m + j} = \sum_{j = 0}^n \binom{n}{j} S_{m + j}

Thus, the final answer would be

2 S_{m + n} + k^{m + n} - \sum_{j = 0}^n \binom{n}{j} S_{m + j} - \sum_{j = 0}^m \binom{m}{j} S_{n + j}\\

So, once you find S_r for all r \le m + n, you can simply find the answer in O(m + n) time since the factorials of binomial coefficient can be precomputed.

If we let

f(x) = \sum_{r = 0}^{\infty} S_r \frac{x^r}{r!} = \sum_{r = 0}^{\infty} \sum_{j = 0}^{k - 1} j^r \frac{x^r}{r!} = \sum_{j = 0}^{k - 1} \sum_{r = 0}^{\infty} \frac{(jx)^r}{r!} = \sum_{j = 0}^{k - 1} e^{jx} = \frac{e^{kx} - e^x}{e^x - 1}.

Thus, if we find the first m + n coefficients in the power series of f(x) we’d be done.

The numerator is easy to deal with since it’s simply

e^{kx} - e^x = \sum_{j = 0}^{\infty} (k^j - 1) \frac{x^j}{j!}

For the denominator, note that for each power series there exists it’s inverse, ie, for each g(x) = \sum_{j = 0}^{\infty} a_j x^j there exists an h(x) = \sum_{j = 0}^{\infty} b_jx^j such that h(x) g(x) = 1.
Therefore, we need to find the inverse power series of

e^x - 1 = \sum_{j = 1}^{\infty} \frac{x^j}{j!}.

However since we only need to find the first m + n coefficients of f(x) we can truncate the sequence at m + n, ie, find the power series

\frac{1}{\sum_{j = 1}^{m + n} \frac{x^j}{j!}}.

after which we can simply multiply the two sequences to get our required S.

Finding the inverse is a classic application of Newton-Ralphson’s method, for a polynomial f let g_j be the polynomial with coefficients such that the first 2^j coefficients for g_jf is 1. Then we get the relation

h(g) = f - g^{-1} \implies g_{n + 1} = g_n - \frac{h(g_n)}{h'(g_n)} = g_n - \frac{f - g_n^{-1}}{g_n^{-2}} = g_n(2 - g_nf)

this is precisely the method described in CF user adamant’s article here. There is also an alternate proof described in CF 438E’s editorial.

The most important thing that has not been described here is to multiply two polynomials efficiently. This is known as FFT and you can read more about it in adamant’s article linked above and the much-revered CLRS.
Note that we need to multiply the polynomials modulo 10^9+7 for which reading the section “Multiplication of arbitrary modulus” on page 6 on adamant’s article will be helpful.
Along with this cp-algorithms.com (translation of e-maxx.ru) has an article on FFT.

Also, while searching around I also found SERSUM - Editorial which was a problem on Long May 2018 Div 1 and the part about computing L in the editorial is exactly the problem about computing S tackled there.
And PFRUIT - Editorial too.
admant also has a blog post on computing S.

The final time complexity of the solution would be O((m + n) \log(m + n)) however, note that you can simply precompute the coefficients of \frac{e^x}{e^x - 1} till \max(m + n) and then multiply k^j to the j-th coefficient to get the sequence \frac{e^{kx}}{e^x - 1} from here, all you need to do is subtract the two series which would give a time complexity of O(a \log a + (m + n) T) where T is the number of test cases and a = \max(m + n) which is 2\times 10^5 here.

I tried to implement my solution at like 7:00 AM but I gave up when I realised when something was wrong with my attempt to find \frac{x}{e^x - 1} since I started getting waaayyy off answers if I tried to find more than ~30 coefficients of the series. Here’s my failed try if you want to see.

tl;dr: GenFun are cool, kids.

6 Likes