[AtCoderやさしい解説]ABC194 C問題 C – Squared Error

\(O(N^2)\)では間に合わない

最初にぱっと思いつくのが\(i\)と\(j\)で二重ループを回して全通り足していくというものだと思います。

しかし制約を見てみると\(2 ≦ N ≦ 2 \times 10^5\)となっています。

明らかに\(O(N^2)\)解法の妨害ですね。

ということで頑張って計算量を削ってみましょう。

式変形をしてみる

こういう時は、同じ計算をしているところを削減したりできることがあるのでそこを発見しやすいように式変形をしてみます。

文字を使うとややこしいので具体的な入力例を使ってみます。

3
1 3 5

計算の順番としては\(a_i\)を左から見ていってそれより右のやつとの差の二乗を足していく感じにします。

よってまず1に着目すると

3 → 5

の順番です。

さらに3に着目すると

5のみ

となります。

5より右には何もないので何もしません。

これらを分けて考えます。

まず1について着目したときの和を\(div_0\)とすると

\[div = (1 – 3)^2 + (1 – 5)^2\]

これを展開して整理すると

\[div = (1^2 + 3^2 – 2 \times 1 \times 3) + (1^2 + 5^2 – 2 \times 1 \times 5)\\
= 2 \times 1^2 -2 \times 1(3 + 5) + (3 ^ 2 + 5 ^ 2)・・・①\]

これらをちょっと観察していきます。

①では\(1^2\)に2がかけられている理由はもちろん1が含まれる計算が二回出てくるからです。

これは着目している数の右にある数の個数と言い換えられます。

よって①での\(1^2\)の係数2は\(n – i – 1\)と表せます。

また①の二項目を見てみると\(2 \times 1(3 + 5)\)となっています。

2は関係ないので\(1(3 + 5)\)を見ます。

この1はいま1に着目しているからでてきたものですよね。

よって①での二項目に出てくる1は\(a_i\)と表せます。

また\((3 + 5)\)は1より右の数の和ですよね。

つまり今見ている数より右の数の和と言い換えられます。

同様に、\((3 ^ 2 + 5 ^ 2)\)は今見ている数より右の数の二乗の和と言い換えられます。

いったん一般化したもので整理すると

\[div = (n – i + 1){a_i}^2 -2a_i(\text{右の数の和}) + (\text{右の数の二乗の和})\]

右の数の和と二乗の和を漸化式で求める

それぞれの\(a_i\)について右の数の和をとっていっていたら結局\(O(N^2)\)かかってしまいます。

ここでこれらを漸化式で求めることを考えてみます。

\(a_i\)の総和を\(sum\)、\(i\)番目より右の数の和をrightとすると

\(i = 0\)のとき

\[right = sum – a_0\]となりますよね。

\(i = 1\)のときは

\[right = right – a_1\]です。

こうやって一回一回一個前の右の数の和からいまの\(a_i\)を引けばいいのです。

さらに\(sum\)はもう使わないので再利用していきます。

よって最初に\(a_i\)の総和を\(sum\)に入れておいて

\[sum = sum – a_i\]としていけばよいのです。

二乗の和についてもまったく同様です。

これで大体の流れは書いたのでプログラムに行きましょう。

ABC194 C問題の私の答え

#include<bits/stdc++.h>
using namespace std;
using lint = long long;
using P = pair<lint, lint>;
const int mod = 1000000007;
#define all(a) (a).begin(),(a).end()
#define rep(i, n) for (lint i = 0; i < (lint)(n); i++)
inline lint fcl(lint num){
    return (num == 1 ? 1 : (num * fcl(num - 1)));
}
int main(){
    lint n;
    cin >> n;
    vector<lint> a(n);
    lint sum = 0;//右の数の和
    lint sumt = 0;//右の数の二乗の和
    rep(i, n){
        cin >> a[i];
        sum += a[i];
        sumt += a[i] * a[i];
    }
    lint ans = 0;
    for(lint i = 0; i < n - 1; i++){
        sum -= a[i];
        sumt -= a[i] * a[i];
        ans += (n - i - 1) * a[i] * a[i];
        ans -= 2 * a[i] * sum;
        ans += sumt;
    }
    cout << ans << endl;
}

こんなかんじですね。

おわり

良ければこの色変記事もどうぞ。

ありがとうございました。