ARC106 D. Powers

問題

長さNの整数列A_iと、整数Kが与えられる。
1 <= X <= Kを満たす整数Xそれぞれについて、
(Σ_{L=1}^{N-1} Σ_{R=L+1}^{N} (A_L + A_R)^X ) mod 998244353
を求めよ。

制約

• 2 <= N <= 2 * 10^5
• 1 <= K <= 300
• 1 <= A_i <= 10^8

解法

Σ_{L=1}^{N-1} Σ_{R=L+1}^{N}という形は、行列で考えるところの、対角を除く上三角部分の部分に相当する。
これは、対称性から、「(行列全体 - 対角部分) / 2」で求められる。
したがって、行列全体「Σ_{L=1}^{N} Σ_{R=1}^{N} (A_L + A_R)^X」と対角部分「Σ_{L=1}^{N} (A_L + A_L)^X」が求められばよい。
対角部分は2^X *A_L^Xなので、前計算で求めておける。

次に、「(A_L + A_R)^X」という形に注目する。
これは、「(x+y)^n」の形なので、二項定理より
Σ_{k=0}^{X} Comb(X,k) * A_L^{X-k} * A_R^{k}
と書くことができる。

「行列全体」の式に入れると、kはL,Rと独立でΣの一番外側に持っていけるので、
Σ_{k=0}^{X} Comb(X,k) Σ_{L=1}^{N} Σ_{R=1}^{N} ( A_L^{X-k} * A_R^{k} )
という形になる。

LとRの部分は、「ΣΣx_i*y_j」の形になっており、これは「(Σx_i) * (Σy_j)」となるため、結局、「ΣA_i^k」の値を前計算で求めておけば、高速に求められる。

前計算はO(NK)、各Xでの計算にO(X)程度なので全体でO(K^2)程度になる。