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;

}