数列の連続部分列の"区間和の2乗"の総和をO(N)で求めるテク
なんだかよくわからないタイトルでごめんなさい。こういうのなんて言ったらいいんだろうか。
先日の Facebook Hacker Cup 2016 Round2 C でそんなテクが要求されて解けなかったので自分用メモ書き。
教えて下さった@kyuridenamidaさんに感謝。
@koyumeishi_ (X-x1)^2+(X-x2)^2+...+(X-x_n)を展開するとnX^2-2*X*(x1+x2+..xn)+(x1^2+x2^2+...+xn^2)になるので,一次の総和と二次の総和持っとくといいかんじです.
— きゅうり (@kyuridenamida) 2016, 1月 23
長さ $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取らず、オーバーフローも気にしてません。