题解：P12716 [Algo Beat Contest 002 C] Counting Square Numbers

dengrunze2608 · 2025-06-07 23:46:49 · 题解

大意

题目需要对数组中的每一个位置 i，找出所有包含 i 的子数组，使得这些子数组的和是一个完全平方数。最后输出满足条件的子数组的个数。

思路

为了方便计算每个子数组的和，我们可以先计算前缀和数组 a，其中 a[i] 表示数组前 i 个元素的和。

对于每一个位置 i，我们可以找到所有包含 i 的子数组，即子数组的起始位置 l 和结束位置 r 满足 l\le i\le r。然后计算这些子数组的和，并检查是否为完全平方数即可。

#include<bits/stdc++.h>
using namespace std;
long long a[5005],ans,sum,s;
int main(){
    int n;
    cin>>n;
    a[0]=0;
    for(int i=1;i<=n;i++){
        cin>>a[i];
        a[i]+=a[i-1];
    }
    for(int i=1;i<=n;i++){
        ans=0;
        for(int l=1;l<=i;l++){
            for(int r=i;r<=n;r++){
                sum=a[r]-a[l-1];
                s=sqrt(sum);
                if(s*s==sum){
                    ans++;
                }
            }
        }
        cout<<ans<<"\n";
    }
    return 0;
}

但是，交上去会发现全部 TLE。 提交记录

计算时间复杂度会发现，因为 n\le5000，所以最坏情况可能为 O(n^3)，绝对会超时。

因此我们需要优化内层循环，可以利用差分进行优化，将时间复杂度由 O(n) 降到 O(1)。

具体来讲，可以开一个差分数组 c，令 c[0]=0 并让 c[i]=a[i]-a[i-1]。这样 c 数组中就是相邻两数的差值了。这样，当求子数组 a[l] 到 a[r] 的和时，这个和实际上等于 c[l]+c[l+1]+...+c[r-1]+c[r]，我们可以发现这个结果恰好等于 a[r]-a[l-1]。

在本题中，我们可使用差分数组 c 来记录每个位置被多少符合条件的子数组覆盖。当发现子数组 [l,r] 的和是完全平方数时，执行 c[l]+=1 以及 c[r+1]-= 1，这样在后续计算 c 数组的前缀和时，位置 l 到 r 会自动加 1。

最后，计算差分数组 c 的前缀和，就能得到每个位置被多少个符合条件的子数组覆盖，得出答案，同时大大降低了时间复杂度。

AC代码

#include<bits/stdc++.h>
using namespace std;
long long a[5005],n,c[5005];
long long sum,s,ans=0; 
signed main(){
    cin>>n;
    a[0]=0;
    for(int i=1;i<=n;i++){
        cin>>a[i];
        a[i]+=a[i-1];
    }
    for(int l=1;l<=n;l++){
        for(int r=l;r<=n;r++){
            sum=a[r]-a[l-1];
            s=sqrt(sum);
            if(s*s==sum){
                c[l]++;
                c[r+1]--;
            }
        }
    }
    for(int i=1;i<=n;++i){
        ans+=c[i];
        cout<<ans<<"\n";
    }
    return 0;
}