D - Powers Editorial by evima

Instead of finding: $\( \displaystyle \sum_{L=1}^{N-1} \sum_{R=L+1}^{N} (A_L+A_R)^X \)$

we will first find:

\[ \displaystyle \sum_{i=1}^{N} \sum_{j=1}^{N} (A_i+A_j)^X \]

and subtract \( \displaystyle \sum_{i=1}^{N} (A_i+A_i)^X \) from it and then divide the result by \(2\).

We have:

\[ \begin{aligned} \displaystyle \sum_{i=1}^{N} \sum_{j=1}^{N} (A_i+A_j)^X &= \sum_{i=1}^{N} \sum_{j=1}^{N} \sum_{k=0}^{X} \binom{X}{k}A_i^k A_j^{X-k} \\ &= \sum_{i=1}^{N} \sum_{j=1}^{N} \sum_{k=0}^{X} X! \cdot \frac{A_i^k}{k!} \cdot \frac{A_j^{X-k}}{(X-k)!} \\ &= X! \sum_{k=0}^{X} \left(\sum_{i=1}^{N} \ \frac{A_i^k}{k!}\right ) \left(\sum_{j=1}^{N} \frac{A_j^{X-k}}{(X-k)!} \right ) \end{aligned} \]

Thus, by pre-computing \( \displaystyle \sum_{i=1}^{N} \ \frac{A_i^k}{k!}\) for every \(k\) in a total of \(O(NK)\) time, we can find \( \displaystyle \sum_{i=1}^{N} \sum_{j=1}^{N} (A_i+A_j)^X \) for each \(X\) in \(O(K)\) time.

The total time complexity of this solution is \(O(NK + K^2)\). (We can also compute the answer in \(O(NK + K \log K )\) time using convolution.)

last update: