CSES: Subarray Sums II

CSES的题都很典，这个就很像是17或以前的 USACO银题目。

这道题简单来说就是给定一个序列​a，找出​a中子序列和为​x的数量。通过前缀和，我们可以将问题变为：给定一个前缀和序列，求​psum[i]-psum[j-1]=x有几对​(i,j)？

思路1

因为数据规模​1\le n\le 2\cdot10^5，不能枚举两个端点。那给​psum排个序，然后跑双指针​l,r使得​psum[r]-psum[l-1]=x。大于就动​l，小于就动​r。

代码1(WA)：

*
 * Created by: Friedforks
 * Date: 2024-11-13
 */
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int n, x;
    cin >> n >> x;
    vector<int> a(n);
    for (int i = 0; i < n; i++)
        cin >> a[i];
    ll ans = 0, acc = 0;
    vector<ll> psum(n + 1);
    for (int i = 0; i < n; i++) {
        psum[i + 1] = psum[i] + a[i];
    }
    sort(psum.begin(), psum.end());
    int l = 0, r = 0;
    for (int i = 0; i <= n; i++) {
        while (psum[i] - psum[l] > x)
            l++;
        if (psum[i] - psum[l] == x && i != l)
            ans++;
    }
    cout << ans << endl;
    return 0;
}

这个想法很轻松就WA了，以下数据有问题：

input

5 0
0 0 0 0 0

correct output:

15

program output:

5

这是因为双指针单向的特性使得​l,r的状态不可能在，比如说​l=0,r=4之后还出现​l=1,r=2的情况。但是这两种在以上数据下都有，所以就少数了。而且还有很多其他的bug，所以这个思路不对。

思路2

参考 Solution - Subarray Sums II (CSES) · USACO Guide

可以使用 map来记录。回顾问题，我们现在假设​psum里已经有一个点在​r，现在需要知道​r前面出现了几次​l，使得​r-l=x，map可以在​\mathcal{O}(log(n))里进行所有操作，记录下前面出现过几次​l：map<l,l count>。 注意！使用 unordered_map会超时，因为他最坏情况下会是​\mathcal{O}(N)，使得程序成为​\mathcal{O}(N^2)。实测最后一个点会超时。

代码2：

/*
 * Created by: Friedforks
 * Date: 2024-11-13
 */
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int n, x;
    cin >> n >> x;
    vector<int> a(n);
    for (int i = 0; i < n; i++)
        cin >> a[i];
    ll ans = 0, psum = 0;
    map<ll, int> cnt;
    cnt[0] = 1;
    for (int i : a) {
        psum += i;
        ans += cnt[psum - x];
        cnt[psum]++;
        cout << "psum: " << psum << "\tcnt[psum-x]" << cnt[psum - x] << endl;
    }
    cout << ans << endl;
    return 0;
}