AtCoder Regular Contest 106 D - Powers

解法

という問題設定はよく登場します。このような問題では、 $l \lt r$ という条件を外して

$1 \le l \le N$ かつ $1 \le r \le N$ である全てのペア $(l, r)$ について、【 $l$ と $r$ について対称な式】の値を合計せよ。

と考えると見通しが良くなることが非常に多いです。この値を求めてから、 $l=r$ の場合を引き、対称性を利用して $2$ で割ると元々の答えが求められることになります。

ということで、 $(A_{l} + A_{r})^{X}$ を $1 \le l \le N$ かつ $1 \le r \le N$ である全てのペア $(l, r)$ について合計することを考えます。

ここで $(A_{l} + A_{r})^{X}$ が

$\displaystyle (A_{l} + A_{r})^{X} = \sum_{k=0}^{X}\left( {}_{X}C_{k}\times{A_{l}}^{k}\times{A_{r}}^{X-k}\right)$

と二項係数を用いて分解できることに注目します。これを全てのペア $(l, r)$ について合計したものが、今求めたい値です。

ここで和を取る順番を逆にします。ある $k$ を1つ固定し、全てのペア $(l, r)$ について、 $A_{l}, A_{r}$ それぞれの次数が $(k, X-k)$ となる項だけを取り出してみましょう。そうするとそれらの合計は、次の形に因数分解できます。

$\displaystyle {}_{X}C_{k}\times({A_{1}}^{k} + \cdots + {A_{N}}^{k})\times({A_{1}}^{X-k} + \cdots + {A_{N}}^{X-k})$

実際にこの式を展開すると、全てのペア $(l, r)$ に対する ${}_{X}C_{k}\times{A_{l}}^{k}\times{A_{r}}^{X-k}$ が全て足された形になっていることが確認できます。

この値を $k$ について合計すれば良いので、今求めたい値は

$\displaystyle\sum_{k=0}^{X}\left({}_{X}C_{k}\times({A_{1}}^{k} + \cdots + {A_{N}}^{k})\times({A_{1}}^{X-k} + \cdots + {A_{N}}^{X-k})\right)$

となります。これは $A_{i}$ の $k$ 乗和を全て前計算しておくことで、 $O(X)$ で求めることができます。

ここから $l=r$ の場合を引き、対称性を利用して $2$ で割ると答えになります。 $l=r$ の場合の総和も、 $A_{i}$ の $X$ 乗和を用いて $2^{X}({A_{1}}^{X} + \cdots + {A_{N}}^{X})$ と計算できます。

計算量は前計算パートが $O(NK)$ 、全ての $X$ に対して答えを求めるパートが $O(K^{2})$ となります。