F - Range Power Sum Editorial
by
spheniscine
Another solution
Let \(\displaystyle dp(r, k) = \sum_{l = 1}^r \left( \sum_{i = l}^r A_i \right)^k\) for \(0 \le r \le N, 0 \le k \le K\). Our desired answer is \(\displaystyle \sum_{r=1}^N dp(r, K)\). Base case: \(dp(0, k) = 0\).
Knowing \(dp(r, *)\), let’s derive a formula for \(dp(r+1, k)\).
\(\displaystyle dp(r + 1, k) = \sum_{l = 1}^{r+1} \left( \sum_{i = l}^{r+1} A_i \right)^k\)
\(\displaystyle =A_{r+1}^k + \sum_{l = 1}^r \left(A_{r+1} + \sum_{i = l}^r A_i \right)^k\)
By the binomial theorem, \(\displaystyle (x+y)^n = \sum_{k=0}^n {n \choose k}x^{n-k}y^k\). Substituting \(\displaystyle \left\langle A_{r+1}, \sum_{i = l}^r A_i, k, j \right\rangle\) for \(\langle x, y, n, k\rangle\) respectively:
\(\displaystyle dp(r + 1, k) =A_{r+1}^k + \sum_{l = 1}^r \sum_{j=0}^k {k \choose j} \cdot A_{r+1}^{k-j} \cdot \left(\sum_{i = l}^r A_i \right)^j \)
\(\displaystyle = A_{r+1}^k + \sum_{j = 0}^k {k \choose j} \cdot A_{r+1}^{k-j} \cdot \sum_{l = 1}^r \left(\sum_{i = l}^r A_i \right)^j \)
[Swap independent summations, then move independent factors according to the distributive law]
\(\displaystyle = A_{r+1}^k + \sum_{j = 0}^k {k \choose j} \cdot A_{r+1}^{k-j} \cdot dp(r, j) \)
We’ve thus found a way to derive \(dp(r + 1, *)\) from \(dp(r, *)\), solving the problem. This should take \(O(NK^2)\) time complexity. It could also theoretically be optimized to \(O(NK \log K)\) via FFT, which would be useful if \(K\) were quite larger.
posted:
last update: