ARC106-D Powers を直感的に解く

今更ですが、ARC106-D Powers の簡単な解説を書きます。

問題リンク：

問題概要：

長さ $N$ の数列 $\{A_i\}$ がある。 $1 \leq X \leq K$ それぞれについて、

$\displaystyle{ \left( \sum_{i \lt j} (A_i+A_j)^X \right) \ {\rm mod} \ 998244353 }$
を求めよ。
$1 \leq N \leq 2 \times 10^5$
$1 \leq K \leq 300$
$1 \leq A_i \leq 10^8$

解法：

対称性から $\displaystyle{\sum_{i \lt j} (A_i+A_j)^X = \frac{1}{2} \sum_{i \neq j} (A_i+A_j)^X}$ なので、これを求めます。（こちらの方が後の計算が楽です。）
まず括弧の中を展開すると、

$\begin{aligned} \sum_{i \neq j} (A_i+A_j)^X &= \sum_{i \neq j} \left(\binom{X}{0} A_i^X + \binom{X}{1} A_i^{X-1} A_j + \cdots + \binom{X}{X} A_j^X \right) \\ &= \binom{X}{0} \sum_{i \neq j} A_i^X + \binom{X}{1} \sum_{i \neq j} A_i^{X-1} A_j + \cdots +\binom{X}{X} \sum_{i \neq j} A_j^X \end{aligned}$

となります。
二項係数は高速に前計算できるので、 $\displaystyle{ \sum_{i \neq j} A_i^k A_j^l \ (0 \leq k,l \leq K)}$ が高速に求まればよいです。

突然ですが、次のような長方形を考えてみましょう。

f:id:null_mn:20201030061131p:plain

すると、 $\displaystyle{ \sum_{i \neq j} A_i^k A_j^l}$ は次の赤い部分の面積に一致します。

f:id:null_mn:20201030061029p:plain

これは次の（緑の面積）-（青の面積）で求まります。

f:id:null_mn:20201030061149p:plain

f:id:null_mn:20201030061141p:plain

緑の面積は $(A_1^k + \cdots + A_N^k)(A_1^l + \cdots A_N^l)$ 、青の面積は $A_1^k A_1^l + \cdots + A_N^k A_N^l$ で求まるので、
$\displaystyle{ \sum_{i \neq j} A_i^k A_j^l = \sum_{i} A_i^k \sum_{i} A_i^l - \sum_{i} A_i^{k+l}}$
が成り立つとわかります。