Wannafly挑戰賽5 A珂朵莉與宇宙 前綴和+枚舉平方數

Wannafly挑戰賽5 A珂朵莉與宇宙 前綴和+枚舉平方數

題目描述

給你一個長爲n的序列a，有n*(n+1)/2個子區間，問這些子區間裏面和爲完全平方數的子區間個數

輸入描述:

第一行一個數n
第二行n個數表示序列a

輸出描述:

輸出一個數表示答案

示例1

輸入
6
0 1 0 9 1 0
輸出
11

備註:

1 <= n <= 100000
0 <= ai <= 10

思路

用前綴和來求任意字段的和
長度爲n的序列a[1],a[2], ….,a[n]
有前綴和prefix_sum[1],prefix_sum[2],….,prefix_sum[n]
若子段a[j],a[j+1],…,a[i] 的和爲完全平方數,則prefix_sum[i]-prefix_sum[j] =square*square 。
對每個prefix_sum[i] 枚舉square的值,若prefix_sum[i]-square*square是一個prefix_sum,則找到一個滿足條件的值。

代碼

#include <bits/stdc++.h>
int is_sum[1000005], a[100005];
int n;
int prefix_sum;
long long ans;
int main(){
    scanf("%d",&n);
    is_sum[0]=1; 
    for(int i=1;i<=n;++i){
        scanf("%d",a+i);
    }
    for(int i=1;i<=n;++i){
        prefix_sum+=a[i];
        for(int square=0;square<=1000 && square*square <= prefix_sum;++square)
            ans+=is_sum[prefix_sum-square*square];
        ++is_sum[prefix_sum];
    }
    printf("%lld\n",ans);
    return 0;
}