atcorder abc194のc問題がtleになる。

Question

atcorderのabc194回のc問題がtleになります。https://atcoder.jp/contests/abc194/tasks/abc194_c
調べてみたところテストケースのsame0,same1,same3が引っかかっているようです。
テストケース

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
template <typename T> bool chmax(T &a, const T &b) { if(a < b) {a = b; return true;} return false; }
template <typename T> bool chmin(T &a, const T &b) { if(a > b) {a = b; return true;} return false; }
const long long  INF=1LL<<60;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)

int main(){
    ll n; cin >> n;
    ll a[401]={0};
    ll ans=0;
    ll numbers[n+1];
    rep(i,n){
        cin >> numbers[i];
    }
    rep(i,n){
        ll b=numbers[i];
        if(a[b+200]){
            ans+=a[b+200];
        }else{
            ll sub=0;
            rep(j,n){
                sub+=(b-numbers[j])*(b-numbers[j]);
            }
            a[b+200]=sub;
            ans+=sub;
        }
    }
    cout << ans/2 <<endl;
}

上が僕の書いたもので、配列a に1度答えを出した物を入れておき、2回目以降はその配列から答えを取り出すようにしています。僕はここの処理時間が長いのではないかと勝手に考えていますがどうなのでしょうか？どこが原因で時間がかかっているのか教えていただきたいです。お願いします！

Answer 1

テストケースの名前からわかると思いますが、same0,same1,same3はすべて同じ数字からなります。
すべて同じ数字だった場合、(A_i - A_j)^2 が常に 0 となるため、一行分足した結果も 0 となります。
つまり、計算結果を a[b+200] に入れてあげても、最初の値の 0 から変わりません。
そうなると、すべての行について、未計算だと勘違いして毎回一行分の足し算が行われることになります。
その結果、O(n^2)の計算が行われて、TLE となっているのだと思われます。

Comments

[y-kou]

回答ありがとうございます。ご教示いただいた通りだったようで、未回答の要素を0から-1にしたところ無事解決しました！