牛客题解 珂朵莉与宇宙

链接：https://ac.nowcoder.com/acm/problem/14600

来源：牛客网

题解

作者 岛田小雅

这道题很仁慈地直接告诉了我们子区间的个数，如果直接暴力遍历所有的子区间，复杂度是 \(O(\frac{n\times (n+1)}{2})\) 那肯定会 TLE。怎么办呢？如果有一种方法能在从左往右遍历数组 \(a\) 的时候就把所有信息都处理好，接下来只需要调用处理好的信息进行计算就好了。

首先，因为要求子区间的和，我们先弄个前缀和数组出来。这样我们就能轻易通过对前缀和的处理来找到需要的区间。对每个 \(qzh_i\)，我们直接对所有的 \(qzh_j\)（\(j<i\)）进行计数，然后遍历所有平方数 \(q_k\)（\(q_k\leqslant qzh_i\)）。在每个 \(qzh_i\) 前出现的 \(qzh_j=qzh_i-q_k\) 的数量之和就是答案。

我们知道，前缀和最大应该是 \(n\times 10=1\times 10^6=N^2\)。也就是说，对每个数遍历比它小的所有平方数，它的复杂度是 \(O(n\times N)\)。总复杂度是 \(O(n+n\times N)\)，问题不大。

注意，因为一共有 \(\frac{n\times (n+1)}{2}\) 个子区间，所以答案要开 \(\texttt{long long}\)（我没注意到所以 WA 了 2 发）。

AC 代码

作者 岛田小雅

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

const int N = 1e5+2;
const int MAXQZH = 1e6+2;
int n;
int a[N], qzh[N];
int cnt[MAXQZH];
long long ans;

int main()
{
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n;
    cnt[0]++; // 因为qzh[i]如果自己就是个平方数的话，它可以选择从自己到前面的所有数作为一个区间，所以一开始就有一个0（相当于qzh[0]合法）
    for(int i = 1; i <= n; i++)
    {
        cin >> a[i];
        qzh[i] = qzh[i-1]+a[i];
        for(int j = 0; j*j <= qzh[i]; j++) ans += cnt[qzh[i]-j*j];
        cnt[qzh[i]]++;
    }
    cout << ans;
    return 0;
}