AtCoder Beginner Contest 399 F問題 Range Power Sum

問題
提出解答
問題の概要
- 制約
解法

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

長さ $N$ の非負整数列 $A = (A_i)$ に対して, 以下を求めよ.

$\displaystyle \sum_{1 \leq l \leq r \leq N} \left(\sum_{i=l}^r A_i \right)^K$

制約

$1 \leq N \leq 2 \times 10^5$
$1 \leq K \leq 10$
$0 \leq A_i \lt 998244353$
入力はすべて整数

解法

$r = 0, 1, 2, \dots, N$ 及び $k = 0, 1, 2, \dots, K$ に対して,

$\displaystyle T_{r, k} := \sum_{l=1}^r \left(\sum_{i=l}^r A_i \right)^K$

とする. このとき, 求めるべき最終解答は

$\displaystyle \sum_{r=1}^N T_K$

である.

このとき, 二項定理を利用することによって,

$\begin{align*} T_{r,k} &= \sum_{l=1}^r \left(\sum_{i=l}^r A_i \right)^K \\ &= \sum_{l=1}^{r-1} \left(\sum_{i=l}^r A_i \right)^K + A_r^K \\ &= \sum_{l=1}^{r-1} \left(\sum_{i=l}^{r-1} A_i + A_r \right)^K + A_r^K \\ &= \sum_{l=1}^{r-1} \sum_{m=0}^k \dbinom{k}{m} A_r^m \left(\sum_{i=l}^{r-1} A_i \right)^{k-m} + A_r^K \\ &= \sum_{m=0}^k \dbinom{k}{m} A_r^m \left(\sum_{l=1}^{r-1} \left(\sum_{i=l}^{r-1} A_i \right)^{k-m} \right) + A_r^K \\ &= \sum_{m=0}^k \dbinom{k}{m} A_r^m T_{r-1, k - m} + A_r^K \\ \end{align*}$