Official

C - Sum of Product Editorial by en_translator


If you evaluate \(A_iA_j\) for all pairs \(i,j\) with \(i < j\), the computational cost will become \(O(N^2)\) which does not finish within the time limit. We need to optimize it somehow.

Out of nowhere, let us define \(S_i=A_1+A_2+\dots+A_i\) for each \(i=1,2,\dots,N\). Then

\(\displaystyle \sum_{1\leq i< j\leq N} A_iA_j = \sum_{j=2}^{N}\sum_{i=1}^{j-1} A_iA_j = \sum_{j=2}^{N}\left(A_j\cdot\left(\sum_{i=1}^{j-1} A_i\right)\right)=\sum_{j=2}^{N}A_jS_{j-1},\)

so if we precompute the values of \(S=(S_1,S_2,\dots,S_N)\), the answer can be found in \(O(N)\) time.

The values of the elements of the sequence \(S\) can be computed for each \(S_1,S_2,\dots,S_N\) in order, using the relations \(S_{i+1}=S_i+A_{i+1}\) to easily obtain \(S_{i+1}\) based on the value \(S_i\), in a total of \(O(N)\) time. (This is an important algorithm called cumulative sums.)

By using the cumulative sums, the problem has been solved in a total of \(O(N)\) time.

As in the sample code below, one can avoid implicitly constructing the sequence \(S\) and compute the answer and \(S_i\) at the same time.

Sample code (C++):

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    int n;
    cin >> n;
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }

    ll ans = 0, sum = 0;
    for (int i = 0; i < n; i++) {
        ans += sum * a[i];
        sum += a[i];
    }

    cout << ans << endl;
}

posted:
last update: