I - Range Power Sum Editorial
by
Iroha_3856
K=10^5でも解ける解法
累積和 \(\displaystyle S_0 = 0, S_i := \sum_{1 \leq j \leq i} A_j (1 \leq i \leq N)\) として、\(\displaystyle \sum_{0 \leq l < r \leq N} (S_r - S_l)^ K\) を求めても答えは変わりません。
平方分割で解きます。
ブロックサイズ \(B\) をとり、\(N+1\) 要素の \(S\) を 順に \(B\) 個ずつ、\(M = \left\lfloor (N+B)/B \right\rfloor\) 個のブロックに分割します。
以下のように二つの問題を考え、それぞれの答えを足し合わせます。
- 問題 $1$:同じブロック内にあるものについて計算する。
ブロック $M$ 個それぞれについて $O(B^2)$ 回 $O(\log K)$ で累乗計算を行うので、計算量は $O(NB \log K)$ となる。 - 問題 $2$ :違うブロック内にあるものについて計算する。
$a^K = K![x^K] \exp(ax)$ を用いれば、 $i$ 番目のブロックに含まれるindexの集合を $D_i \ (1 \leq i \leq M)$ として、答えは $$\displaystyle \sum_{1 \leq i < j \leq M} \sum_{l \in D_i} \sum_{r \in D_j} (S_r - S_l)^K$$ $$=\displaystyle K!\sum_{1 \leq i < j \leq M} [x^K] \sum_{l \in D_i} \sum_{r \in D_j} \exp((S_r - S_l)x)$$ $$= \displaystyle K! \sum_{1 \leq i < j \leq M} [x^K] \sum_{l \in D_i} \sum_{r \in D_j} \exp(-S_l x)\exp(S_r x)$$ $$= \displaystyle K!\sum_{1 \leq i < j \leq M} [x^K] (\sum_{l \in D_i} \exp(-S_l x))(\sum_{r \in D_j} \exp(S_r x))$$ すなわち、$F_i = \sum_{l \in D_i} \exp(-S_l x), G_i = \sum_{r \in D_j} \exp(S_r x)$ と定義すると、答えは $$\displaystyle K!\sum_{1 \leq i < j \leq M} [x^K] F_i G_j$$ です。
$F_i, G_i$ をそれぞれ求めることを考えましょう。 $$\displaystyle F_i = \sum_{k = 0}^\infty \frac{\sum_{t \in D_i} (-S_t x)^k}{k!} = \sum_{k = 0}^\infty (\sum_{t \in D_i} (S_t)^k) \frac{(-1)^k}{k!} x^k $$ $$\displaystyle G_i = \sum_{k = 0}^\infty \frac{\sum_{t \in D_i} (S_t x)^k}{k!} = \sum_{k = 0}^\infty (\sum_{t \in D_i} (S_t)^k)\frac{1}{k!} x^k $$ であるので、$\displaystyle P_k = \sum_{t \in D_i} (S_t)^k$ を $O(B \log^2 B + K \log K)$ で計算することで、$F_i, G_i$ は $O(K)$ 時間で求まります。$i$ 乗和の列挙と呼ばれる有名問題(多項式・形式的べき級数-高速に計算できるもの - $i$ 乗和の $i$ に関する列挙)です。 $O(B \log^2 B)$ で $\prod (1 - A_i x)$ が求まり、この $\log$ を取るのに $O(K \log K)$ 時間かかります。
$M$ 個のブロックすべてについてこれを行うので、$O(MB \log^2 B + MK \log K) = O(N \log^2 B + NK/B \log K)$ 時間で $F, G$ を計算できます。
$F, G$ が求まったとき、$\displaystyle \sum_{1 \leq i < j \leq M} [x^K] F_i G_j = \sum_{1 \leq j \leq M} [x^k] G_j (\sum_{1 \leq i < j} F_i)$ であるので、$j$ を昇順に動かしつつ、 $\displaystyle \sum_{1 \leq i < j} F_i$ を各 $j$ について $O(K)$ 時間で差分更新しながら、$[x^K]$ も $O(K)$ 時間で求ていけば、このパートは $O(KM) = O(NK/B)$ 時間で行えます。 $F, G$ を列挙するパートがボトルネックとなって、問題 $2$ にかかるのは $O(N \log^2 B + NK/B \log K)$ 時間です。
以上より、問題 \(1\) と \(2\) を合わせると \(O(N \log^2 B + NK/B \log K + NB \log K)\) 時間でこの問題を解くことができます。
\(B = \sqrt{K}\) としてとれば、計算量は \(O(N \log^2 K + N \sqrt{K}\log K)\) 時間となります。
AtCoderのコードテスト上でこのコードを \(N = K = 10^5\) 、\(A\) はランダム、という条件で動かしたところ、\(5300\mathrm{ms}\) 程度で動作しました。
posted:
last update:
