数列の連続部分列の"区間和の2乗"の総和をO(N)で求めるテク

なんだかよくわからないタイトルでごめんなさい。こういうのなんて言ったらいいんだろうか。

先日の Facebook Hacker Cup 2016 Round2 C でそんなテクが要求されて解けなかったので自分用メモ書き。
教えて下さった@kyuridenamidaさんに感謝。

長さ $n$ の 数列 $A_1, A_2, ... , A_n$ について、 区間 $[l,r]$ の和を $ S_{l,r} = A_l + A_{l+1} + ... + A_r$ とします。 全ての区間について、${S_{l,r}}^2$ の総和を求めます。
$$ T(n) = \sum_{l=1}^{n} \sum_{r=l}^{n} {S_{l,r}}^2 $$ 例えば数列 $[1,2,3]$ は、$1^2 + 2^2 + 3^2 + (1+2)^2 + (2+3)^2 + (1+2+3)^2 = 84$ となります。

きゅうりさんのツイートによると

$$ (X-x_1)^2+(X-x_2)^2+...+(X-x_n)^2 = nX^2-2X(x_1+x_2+..x_n)+(x_1^2+x_2^2+...+x_n^2) $$

だからなんか嬉しいそうです。この式がなんで嬉しいのかすぐにわからなかったので以下メモ。

まずは効率悪いけど何やってるかわかりやすい$O(N^2)$,$O(N^3)$。


$O(N^3)$

式の通りの3重ループを書けば$O(N^3)$で計算可能です。


$O(N^2)$

累積和を持って置けば$S_{l,r}$が$O(1)$で計算できるので$O(N^2)$でできます。

$A_1, A_2, ... , A_k$ の尻尾に新しく $A_{k+1}$ を追加すると考えて、 右端からの部分和 $S_{x,k+1}$ を更新しながら $O(N^2)$ もアリです。


$O(N)$

これが本命。
$O(N^2)$ の2つ目の手法のように、 $A_{k+1}$ を末尾に追加しながら計算することを考えます。
$$T(k) = \sum_{l=1}^{k} \sum_{r=l}^{k} {S_{l,r}}^2$$ が既に計算出来ているとき、ここに $A_{k+1}$ を追加すると、増分は

$$ \begin{aligned} &(A_{k+1} + A_k + ... + A_3 + A_2 + A_1)^2 + \\ &(A_{k+1} + A_k + ... + A_3 + A_2)^2 + \\ &(A_{k+1} + A_k + ... + A_3)^2 + \\ &\vdots \\ &(A_{k+1} + A_k)^2 + \\ &(A_{k+1})^2 \end{aligned} $$

です。
これよく見ると、きゅうりさんのツイートにあった $(X-x_1)^2+(X-x_2)^2+...+(X-x_n)^2$ の形をしています。

$$ \begin{aligned} &(A_{k+1} + A_k + ... + A_3 + A_2 + A_1)^2 + \\ &(A_{k+1} + A_k + ... + A_3 + A_2 + \color{red}{A_1})^2 + \\ &(A_{k+1} + A_k + ... + A_3 + \color{red}{A_2 + A_1})^2 + \\ &\vdots \\ &(A_{k+1} + A_k + \color{red}{ ... + A_3 + A_2 + A_1 })^2 + \\ &(A_{k+1} + \color{red}{A_k + ... + A_3 + A_2 + A_1 })^2 + \\ \end{aligned} $$

赤色の部分が $A_{k+1} + A_k + ... + A_3 + A_2 + A_1$ から引かれていると考えると、 きゅうりさんの式の変数は次のようになります。

$$ \begin{aligned} X &= A_{k+1} + A_k + ... + A_2 + A_1 &= S_{1,k+1} \\ x_1 &= 0 &= S_{1,0} \\ x_2 &= A_1 &= S_{1,1} \\ x_3 &= A_1+A_2 &= S_{1,2} \\ &\vdots & \\ x_{k+1} &= A_1+A_2+...+A_k &= S_{1,k} \end{aligned} $$

$ X $ , $ \sum x_i $ , $ \sum x_i^2 $ を持ちながら先頭から順に見ていけば $O(N)$ で済みます。コード見た方が速いかも。


素数3000でideoneしてみた結果https://ideone.com/2zKyHZ

計算量 時間(ms)
O(N) 0ms
O(N2) 15ms
O(N2) 累積和 18ms
O(N3) 4159ms

素数20000https://ideone.com/CLM6nG

計算量 時間(ms)
O(N) 0ms
O(N2) 666ms
O(N2) 累積和 807ms
O(N3) 無理

まぁ計算量的に試さなくてもわかりますよね。 実際使わされる場面ではmod 1e9+7とか取らされたりすると思いますが、この実験では特にmod取らず、オーバーフローも気にしてません。