AtCoder Beginner Contest 194 C - Squared Error题解

比赛地址

C - Squared Error

  本题需要我们求出对于\(\displaystyle \sum_{i = 2}^{N} \sum_{j = 1}^{i - 1} (A_i - A_j)^2\)的所有值。由于\(N\)的数据范围为\(2 \le N \le 3 \times 10^5\)，最暴力的算法的时间复杂度为\(O(N^2)\)。因此我们就需要对计算进行优化。

 解法一(从小情况入手)

  我们考虑一下\(N=3\)的情况，现有数组元素\(A_1,A_2,A_3\)，得到的结果就是

\[\begin{align*} res&=(A_2-A_1)^2+(A_3-A_1)^2+(A_3-A_2)^2\\ &=2\cdot A_1^2+2\cdot A_2^2+2\cdot A_3^2-2\cdot A_1\cdot A_2-2\cdot A_1\cdot A_3-2\cdot A_2\cdot A_3\\ &=2\cdot A_1^2+2\cdot A_2^2+2\cdot A_3^2-[A_1\cdot (A_2+A_3)+A_2\cdot (A_1+A_3)+A_3\cdot (A_1+A_2)] \end{align*} \]

我们发现，假设\(tot=A_1+A_2+A_3\),减号后的结果就是\(\sum\limits_{i=1}^nA_i\cdot (tot-A_i)\)。那么面前平方项的常数是什么呢？我们继续增加数组元素的时候，前\(N-1\)个数的平方数都会增加\(1\)，与第\(N\)个数的平方的个数一致，即\(N-1\)个。
  因此我们就推出新的计算公式为\((n-1)\cdot \sum\limits_{i=1}^nA_i^2-\sum\limits_{i=1}^nA_i\cdot (tot-A_i)\)。就可以在\(O(N)\)的时间内算出结果。

参考代码

#include<bits/stdc++.h>

using namespace std;

typedef long long LL;

int main()
{ 
    int n;
    cin>>n;
    vector<LL> a(n);

    LL res=0,tot=0;
    for(int i=0;i<n;i++){
        cin>>a[i];
        res+=a[i]*a[i];
        tot+=a[i];
    }

    LL tmp=0;
    for(int i=0;i<n;i++) tmp-=a[i]*(tot-a[i]);
    
    cout<<res*(n-1)+tmp<<endl;

    return 0;
}

 解法二(参考自官方题解)转化

  我们算法中的性能瓶颈在于\(N\)的数据范围。注意到\(a_i\)的数据范围较小，可以从这里入手。由于数组元素间的差值较小，因此我们可以一次将相同差值计算出。
  我们定义一个数组\(cnt\)，\(cnt[i]\)记录的是数值\(i\)在原数组中出现的次数，由于存在负数，我们可以选择每个数值增加相同的值，比如\(200\)。或者我们可以使用一个哈希表来记录即可。接着，对于\(i\)和\(j (-200 \le i \lt j \le 200)\)，它们的差值出现的次数就是\(cnt[i]\cdot cnt[j]\)，因此计算总值就是\(cnt[i]\cdot cnt[j]\cdot \mid i-j\mid ^2\)，由于相同大小的元素差值为\(0\)，对结果并没有影响，因此不需要计算。

参考代码

#include<bits/stdc++.h>
typedef long long LL;
using namespace std;

int main()
{
    int n;
    cin>>n;
    vector<LL> a(n);
    unordered_map<LL,LL> cnt;
    for(int i=0;i<n;i++){
        cin>>a[i];
        cnt[a[i]]++;
    }
    LL res=0;
    for(int j=-200;j<=200;j++){
        for(int i=-200;i<j;i++){
            res+=cnt[i]*cnt[j]*(i-j)*(i-j);
        }
    }
    cout<<res;
    return 0;
}

 解法三(参考自官方题解)

  我们可以考虑对计算公式直接进行等价转换。

\[\begin{align*} &\sum_{i = 2}^{N} \sum_{j = 1}^{i - 1} (A_i - A_j)^2\\ &=\frac{1}{2}\sum_{i=1}^{N} \sum_{j=1}^{N}(A_i - A_j)^2\\ &=\frac{1}{2}\sum_{i=1}^{N} \sum_{j=1}^{N}(A_i^2 + A_j^2-2\cdot A_i\cdot A_j)\\ &=\frac{1}{2}\sum_{i=1}^{N}(\sum_{j=1}^{N}A_i^2+\sum_{j=1}^{N}A_j^2-2\cdot \sum_{j=1}^{N}A_i\cdot A_j)\\ &=\frac{1}{2}N\cdot\sum_{i=1}^{N} A_i^2+\frac{1}{2}N\cdot\sum_{j=1}^{N}A_j^2-\sum_{i=1}^{N}(A_i\cdot (\sum_{j=1}^{N}A_j))\\ &=N\cdot \sum_{i=1}^{N}A_i^2-(\sum_{i=1}^{N}A_i)^2 \end{align*} \]

至此，本算法就可以在\(O(N)\)的时间内解决本题。

参考代码

#include<bits/stdc++.h>

using namespace std;
typedef long long LL;

int main()
{

    LL res=0,sum=0;
    int n;
    cin>>n;
    vector<int> a(n);
    for(int i=0;i<n;i++){
        cin>>a[i];
        res+=a[i]*a[i];
        sum+=a[i];
    }

    cout<<n*res-sum*sum<<endl;
    
    return 0;
}