久々に AtCoder の解説記事を書こうと思います．今日 AtCoder の読み方がアットコーダーだと知りました．今までエーティーコーダーだと思ってました．
問題
公式問題はコチラ．
問題文
長さ $N$ の数列 $A$ に対して，各要素同士の差の二乗和を求めよ．
$$\sum^N_{i=2} \sum^{i-1}_{j=1} (A_i – A_j)^2$$
制約
- $2 \leq N \leq 3 \times 10^5$
- $|A_i| \leq 200$
- 入力に含まれる値は全て整数
解説
考え方
$0 \leq i, j \leq N, i \neq j$ における数列 $(A_i, A_j)$ のすべての組み合わせに対して差の二乗を求め，これを合計するという方法をためしてみましたが，TLE となってしまいました．
$_n C _2 = \frac{N(N-1)}{2} = O(N^2)$ 回計算が行われるためですね．
公式解説の解説2では，式変形によって計算量を $O(N)$ とする方法が書かれていました．
$$\sum^N_{i=2} \sum^{i-1}_{j=1} (A_i – A_j)^2$$
この式は $i = j$ で $A_i – A_j = 0$ であり，$A_i – A_j = A_j – A_i$ であるため，以下のように変形することができます．
$$\frac{1}{2} \sum^N_{i=1} \sum^{N}_{j=1} (A_i – A_j)^2 = \frac{1}{2} \sum^N_{i=1} \sum^{N}_{j=1} (A_i^2 + A_j^2 – 2A_i A_j)$$
さらに式が簡単になるように変形していくと，以下のようになります．
$$ = \frac{1}{2}N \sum^N_{i=1} A_i^2 + \frac{1}{2}N \sum^{N}_{j=1}A_j^2 \ – \sum^N_{i=1} A_i \sum^{N}_{j=1} A_j $$
$$ = N\sum^N_{i=1} A_i^2 \ – (\sum^N_{i=1} A_i)^2$$
上記の式に書き直せば計算量が $O(N)$ になり，計算が簡略化されます！
ソースコード
#include <bits/stdc++.h>
#include <math.h>
#include <string>
using namespace std;
int main(void) {
int N;
cin >> N;
vector<int> A(N);
for (int i=0; i<N; i++) {
cin >> A[i];
}
long long ans = 0;
long long a1 = 0;
long long a2 = 0;
for (int i=0; i<N; i++) {
a1 += A[i]*A[i];
a2 += A[i];
}
ans = N*a1 - a2*a2;
cout << ans << endl;
return 0;
}まとめ
今回は式変形により計算量を小さくできるよという例でした．式変形を使えば，シグマの数を2から1にできるというのは結構面白いです．この着想は競プロならではかもしれません．
コメント