F - Cumulative Sum Editorial by en_translator
- Preface
- Observation
- Formula of sum of \(K\)-th powers
- Properties of \(f(n, m)\)
- Lagrange interpolation
- Evaluation of \(f(N)\) using Lagrange interpolation
- Summary
- Bonus
1. Preface
This problem treats techniques that are frequently seen as a subproblems of high-level combinatory problems. We hope the knowledge obtained throughout this problem will be a footstep for getting AC for more advanced combinatory problems in the future.
This problem asks to find the \(n\)-th term of a sequence. When it comes to problems that requires faster computation than linear time, it is crucial to understand the properties of the sequence. Shown in the table below is a summary of properties and computational complexity of the sequences that have commonly been seen in the past ABCs.
| type of sequence | example | precalculation | query |
| powers | \(n ^ k\) | None | \(\mathrm{O}(\log k)\) |
| factorials / binomial coefficients | \(n!,\binom{n}{k}\) | \(\mathrm{O}(n)\) | \(\mathrm{O}(1)\) |
| sequences whose transitions are expressed as a \(d \times d\) matrix |
Fibonacci sequence | None | \(\mathrm{O}(d^3 \log k)\) |
Now, what kind of property does the sequence in the problem satisfies?
2. Observation
When you are stuck in such kind of problem, one of the strategy is to observe the behaviors for the simple cases and find good properties.
Let us consider the cases with small \(m\). First of all, \(m = 0\) is trivial; it can be computed easily by
\[ f(n, 1) = n^k.\]
Next, consider the case where \(m = 1\); we have
\[ \begin{aligned} f(1, 1) &=& 1^K \\ f(2, 1) &=& 1^K + 2^K \\ f(3, 1) &=& 1^K + 2^K + 3^K \\ \vdots \end{aligned} \]
so we can figure out that \(f(n, 1)\) (which we will denote \(S_K(N)\) for convenience) can be expressed as a sum of \(K\)-th powers, namely
\[ S_K(n) = \sum_{i = 1} ^ n i^K. \]
Especially for \(K = 1, 2\) and \(3\), the value \(S_K(n)\) can be expressed as a well-known elementary formula, that is,
- If \(K = 1\), \(S_1(n) = \frac{n(n + 1)}{2} \);
- If \(K = 2\), \(S_2(n) = \frac{n(n+1)(2n+1)}{6} \);
- If \(K = 3\), \(S_3(n) = \frac{n^2 (n+1)^2}{4} \).
Observing the formula above, we can propose an important hypothesis that \(S_K(n)\) is a polynomial of degree \(K + 1\). We will prove this hypothesis in the next chapter.
3. Formula of sum of \(K\)-th powers
We will prove the following proposition.
\(\displaystyle S_K(N) = \sum_{n=1}^N n^k\) can be expressed as a polynomial of \(K + 1\) degrees.
There are several ways to prove this. Here we will show it recursively.
Define a function
\[T_K(n)=n^K - (n-1)^K.\]
Here \(T_{K+1}\) can be expressed in the two ways:
\[\displaystyle \sum_{n=0}^N T_{K+1}(n) = N^{K+1},\]
and at the same time
\[\displaystyle \sum_{n=0}^N T_{K+1}(n) = \sum_{n=0}^N Kn^K + (\text{polynomial in }n\text{ of degree }(K-1)).\]
If the proposition holds for every integer less than or equal to \(K\), we have
\[\sum_{n=0}^N T_{K+1}(n) = (\text{polynomial in }n\text{ of degree }K) + K \sum_{n=0}^N n^K,\]
so
\[S_K(N) = \sum_{n=0}^N n^K = \frac{N^{K+1} + (\text{polynomial in }n\text{ of degree }K)}{K}\]
follows. Accompanied by the case \(K=1\), the proposition has been proved recursively.
4. Properties of \(f(n, m)\)
Now let’s get back to the original problem. By the fact shown in Chapter 3, the following property of \(f(N, M)\) can be proved:
When \(M\) is fixed, \(f(N,M)\) can be expressed as a polynomial in \(N\) of degree \(M + K\).
This fact can be shown recursively. Assuming that \(f(N,m)\) can be expressed as a polynomial in \(N\) of degree \(m+K\), due to the fact obtained in Chapter 3 it follows that
\[ \begin{aligned} f(N, m+1) &=& \sum_{n = 1} ^ N f(n, m) \\ &=& \sum_{n=1}^N (\text{polynomial in }n\text{ of degree }(m+K)) \\ &=& (\text{polynomial in }N\text{ of degree }(m+K+1)), \end{aligned} \]
so, accompanied by \(f(n,0)=n^K\) for \(m=0\), it has been proved.
Therefore, we can now see that the desired sequence can be expressed as a polynomial of \(N\) for fixed \(M\). However here comes another issue: how can \(f(N)\) be computed fast when \(f(n)\) is known to be a polynomial?
If we know the coefficients of the polynomial, it can be computed in an \(O(N)\) time with a naive algorithms like Horner’s method. However, in this problem it seems difficult to reconstruct the coefficients of \(f(N, M)\).
In the following chapters, we will explain some properties of polynomials and explain the algorithm to calculate \(f(N)\) fast.
5. Lagrange interpolation
In fact, a polynomial has the following property.
When \(x_0, x_1, \dots,x_d\) and \(y_0, \dots,y_d\) is given (where \(x_i\) are pairwise distinct), a polynomial \(f(x)\) of degree no more than \(d\) such that
\[f(x_i) = y_i (i=0,\ldots,d)\]
is uniquely determined.
Let take linear functions as an example: “A polynomial with degree no more than \(1\) such that \(f(0) = b\) and \(f(1) = a + b\)” is uniquely determined as \(f(x) = ax + b\). The proposition above is its generalization.
The proposition implies the following important fact:
- Given a polynomial \(f(N)\) of degree \(d\) with unknown coefficients, we want to compute \(f(N)\). Here, even if we do not know the coefficients, one can reconstruct \(f(N)\) as long as we know \(f(0), f(1), \dots, f(d)\).
Look back on the original problem; \(f(0, M), f(1, M), \dots, f(d, M)\) can be computed easily with cumulative sums. In such problems, it is essential to hold a polynomial as a form of \(f(0), f(1),\dots, f(d)\) in order to reduce the computational complexity.
Here, the proposition can be solved with the help of a matrix. Let
\[f(x) = f_0 + f_1 x + f_2 x^2 + \dots + f_d x^d,\]
then \((f_0, f_1,\dots,f_d)\) and \((y_0,\dots,y_d)\) can be expressed with the aid of Vandermonde matrix:
\[ \left( \begin{array}{ccccc} 1 & x_0 & x_0^2 & \cdots & x_0^d \\ 1 & x_1 & x_1^2 & \cdots & x_1^d \\ 1 & x_0 & x_0^2 & \cdots & x_0^d \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & x_d & x_d^2 & \cdots & x_d^d \\ \end{array} \right) \left( \begin{array}{c} f_0 \\ f_1 \\ f_2 \\ \vdots \\ f_d \\ \end{array} \right) = \left( \begin{array}{c} y_0 \\ y_1 \\ y_2 \\ \vdots \\ y_d \\ \end{array} \right). \]
Since the determinant of Vandermonde matrix is
\[ \left \vert \begin{array}{ccccc} 1 & x_0 & x_0^2 & \cdots & x_0^d \\ 1 & x_1 & x_1^2 & \cdots & x_1^d \\ 1 & x_0 & x_0^2 & \cdots & x_0^d \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & x_d & x_d^2 & \cdots & x_d^d \\ \end{array} \right \vert = \prod_{j \lt i} (x_i - x_j) \neq 0, \]
we have shown that one can uniquely determine \(f_0,\ldots,f_d\) that satisfies the conditions.
Thus, we have seen that a polynomial \(f(x)\) of degree \(d\) is uniquely determined when \(f(0),f(1),\ldots,f(d)\) are given. However, multiplying the Vandermonde matrix to the vector of \(y\) to obtain \(f(N)\) costs too much, so we introduce polynomials called Lagrange basis polynomials.
Let us define Lagrange basis polynomial \(L_i (x)\) for given \(x_0,\dots,x_d\) as follows:
\[ \begin{aligned} L_i(x) &=& \prod_{j \neq i} \frac{x - x_j}{x_i - x_j} \\ &=& \frac{(x - x_0) (x - x_1)\dots(x-x_{i-1})(x-x_{i+1})\dots (x-x_d)}{(x_i - x_0) (x_i - x_1)\dots(x_i-x_{i-1})(x_i-x_{i+1})\dots (x_i-x_d)} \end{aligned} \]
Then, using \(L_i(x)\), \(f(x)\) is expressed as follows:
Given \(x_0, x_1, \dots,x_d\) and \(y_0, \dots,y_d\), the polynomial \(f(x)\) of degree no more than \(d\) satisfying
\[f(x_i) = y_i (i=0,\ldots,d)\]
can be expressed as follows, using Lagrange basis polynomials \(L_i(x)\):
\[f(x) = y_0 L_0(x) + y_1 L_1(x) + \dots + y_d L_d(x)\]
We will justify the claim above. A Lagrange basis polynomial
\[ L_i(x) = \frac{(x - x_0) (x - x_1)\dots(x-x_{i-1})(x-x_{i+1})\dots (x-x_d)}{(x_i - x_0) (x_i - x_1)\dots(x_i-x_{i-1})(x_i-x_{i+1})\dots (x_i-x_d)} \]
satisfies the following property:
\[L_i(x_j) = \begin{cases} 1 & i = j \\ 0 & i \neq j \end{cases}. \]
By the equation above, it follows that
\[ y_0 L_0(x_i) + y_1 L_1(x_i) + \dots + y_d L_d(x_i) = y_i,\]
so the claim is justified, with the fact that the degree of \(L_i(x)\) is at most \(d\).
Therefore, we have shown that given multiple points \((x, y)\), the polynomial of minimum degree that passes through all the given points is uniquely determined. A discovery of a polynomial that passes through all the given points is called interpolation, and particularly an interpolation with the help of Lagrange basis polynomials is called Lagrange interpolation.
6. Evaluation of \(f(N)\) using Lagrange interpolation
Now that all the preparation is set, we will explain a fast algorithm computing \(f(N)\) from \(f(0),f(1),\dots,f(d)\).
For a polynomial \(f(x)\) of degree at most \(d\), where \(y_0 = f(0),y_1 = f(1) \ldots,\) and \(y_d = f(d)\) are known, \(f(N)\) can be computed in a total of \(\mathrm{O}(d)\) time.
Using the Lagrange basis polynomials for \(x_i = i\), we obtain
\[f(N) = y_0 L_0(N) + y_1 L_1(N) + \dots + y_d L_d(N),\]
so all we need is to compute \(L_i(N)\) fast enough. \(L_i(N)\) can be transformed as
\[ \begin{aligned} L_i(N) &=& \frac {(N - 0) (N - 1)\dots(N-(i-1))(N-(i+1))\dots (N-d)} {(i - 0) (i - 1)\dots(i-(i-1))(i-(i+1))\dots (i-d)} \\ &=& \frac {\displaystyle \prod_{0 \leq j \leq d, j \neq i} (N - j)} {i! \cdot (-1)^{d - i} (d - i)!}. \end{aligned} \]
If we precalculate the cumulative products of \(N-i\) from the left and from the right and \((1!)^{-1},(2!)^{-1},\dots,(d!)^{-1}\) in an \(\mathrm{O}(d)\) time, \(L_i(N)\) can be computed in an \(\mathrm{O}(1)\) time. Therefore \(L_i(N)\) has been computed in a total of \(\mathrm{O}(d)\), and \(f(N)\) has been found in a total of \(\mathrm{O}(d)\) time too.
7. Summary
By the observations above, we can see that the problem can be solved in a total of \(\mathrm{O}(K \log K + (K + M)M)\) with the following simple algorithm:
- find the first \(K+M+1\) terms of \(f(n, 0) = n^K\) naively;
- calculate the cumulative sums \(M\) times to find the first \(K+M+1\) terms of \(f(n, M)\);
- find \(f(N, M)\) with Lagrange’s interpolation.
The implementation is as follows. Since the time limit is a bit strict, acquiring AC with PyPy may require careful implementation such as optimization of division by JIT compile and avoiding bottlenecks of big integers, or it may lead to TLE.
mod = 1000000007
def fast_pow(x, k):
res = 1
while k:
if k & 1:
res = res * x % mod
x = x * x % mod
k >>= 1
return res
N, M, K = map(int, input().split())
S = M + K
dp = [fast_pow(i, K) for i in range(S + 1)]
for _ in range(M):
for i in range(1, S + 1):
dp[i] = (dp[i - 1] + dp[i]) % mod
invfac = [0] * (S + 1)
fac = 1
for i in range(1, S + 1):
fac = fac * i % mod
invfac[-1] = pow(fac, mod - 2, mod)
for i in reversed(range(S)):
invfac[i] = invfac[i + 1] * (i + 1) % mod
N %= mod
L = [1] * (S + 1)
for i in range(S):
L[i + 1] = L[i] * (N - i) % mod
R = [1] * (S + 1)
for i in reversed(range(1, S + 1)):
R[i - 1] = R[i] * (N - i) % mod
answer = 0
for i in range(S + 1):
tmp = dp[i] * L[i] % mod * R[i] % mod * invfac[i] % mod * invfac[S - i] % mod
if S % 2 == i % 2:
answer += tmp
else:
answer -= tmp
print(answer % mod)
8. Bonus
Although we did not calculate the coefficients sequence of the polynomial \((f_0, f_1,\dots,f_d)\) because we did not need them, if we do, it is known that \((f_0, f_1,\dots,f_d)\) can be obtained from \((y_0,\dots,y_d)\) in a total of \(\mathrm{O}(d \log^2 d)\) time with Vandermonde matrix. The Polynomial Interpolation problem in Library Checker has a sample code, so take a look at it if you are interested in.
For \(N\leq 10^{18}, M \leq 10^7\), and \(K \leq 10^6\), find \(f(N, M) \bmod 998244353\).
posted:
last update: