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;
}