C - Sum of Product 解説 by seekworser


\( \displaystyle \sum_{i=1}^N \sum_{j=1}^N A_i A_j \) という式を考えます。これは、\(i\)\(j\) ごとに分けたあと、添字を取り替えることで

\[ \sum_{i=1}^N \sum_{j=1}^N A_i A_j = \left( \sum_{i=1}^N A_i \right) \left( \sum_{j=1}^N A_j \right) = \left( \sum_{i=1}^N A_i \right)^2 \]

のように変形できます。

ところで、\( \displaystyle \sum_{i=1}^N \sum_{j=1}^N A_i A_j \)\(i,j\) の大小関係に着目して分解すると

\[ \sum_{i=1}^N \sum_{j=1}^N A_i A_j = \sum_{1 \le i < j \le N} A_i A_j + \sum_{1 \le i = j \le N} A_i A_j + \sum_{1 \le j < i \le N} A_i A_j \]

のように変形できますが、対称性から第一項と第三項が等しいことに着目し、さらに第二項の添字を取り替えることで

\[ \sum_{i=1}^N \sum_{j=1}^N A_i A_j = 2 \sum_{1 \le i < j \le N} A_i A_j + \sum_{i=1}^N A_i^2 \]

とすることができます。移行して最初の式変形と合わせることで結局

\[ \sum_{1 \le i < j \le N} A_i A_j = \frac{\left( \sum_{i=1}^N A_i \right)^2 - \sum_{i=1}^N A_i^2 }{2} \]

が得られます。右辺の2つの総和はいずれも \(O(N)\) で計算可能なので、この問題を \(O(N)\) で解くことができます。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

int main() {
    int n; cin >> n;
    vector<ll> a(n);
    for (int i=0; i<n; i++) cin >> a[i];
    ll sum = 0, square_sum = 0;
    for (int i=0; i<n; i++) {
        sum += a[i];
        square_sum += a[i] * a[i];
    }
    ll ans = (sum * sum - square_sum) / 2;
    cout << ans << "\n";
}

投稿日時:
最終更新: