Codeforces Round #442 (Div. 2) 877 F - Ann and Books 莫队算法 离散化

题目链接: F - Ann and Books

题目大意

n本书的序列, 第i本书有a[i]若干道数学或者经济学题目, q个询问[li, ri], 要求输出[li, ri]中, 满足”数学题总数比经济学题总数多k道”的子区间数量
数据范围: $1 \leq n \leq 10^5, 0 \leq a[i] \leq 10^9, -10^9 \leq k \leq 10^9, 1 \leq q \leq 10^5$

思路

首先对数据进行预处理
要求数学题数-经济学题数 = k, 那么对经济学题目数取负数(a[i] = -a[i]), 然后求前缀和pre[i], 这样处理后, 子区间[a, b]满足”数学题总数比经济学题总数多k道” 就相当于 pre[b] - pre[a-1] == k
如果我们已经知道了区间[a, b]的答案ans, 并且已经统计好了区间[a-1, b]中所有pre[i]各个数字的个数cnt[pre[i]]

(至于为什么求解[a, b]时要统计[a-1, b]中所有数字的个数, 这是因为判断[x, y]是否满足条件时, 要判断pre[y] - pre[x-1] == k, 需要用到pre[x-1], 所以要统计[a-1, b])

那么对于区间[a, b-1], [a, b+1], [a-1, b], [a+1, b]的答案, 我们可以根据[a, b]的答案$O(1)$求出, 具体看代码, 接下来就可以用莫队算法了

因为pre[i]的范围比较大, 所以需要离散化, 但因为莫队算法的复杂度是$O(q\sqrt{n}) = 3\cdot10^7$, 所以我们获取某个值离散化后的值时, 必须在$O(1)$内求出, 所以不能每次调用都调用map或者lowe_bound(), 需要将所有值的离散化后的值放到一个数组里面, 与原来的值一一对应, 这样就可以$O(1)$得到离散化后的值

代码

GNU C++14 Accepted 420 ms 70600 KB

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#include <map>
#include <set>
#include <cmath>

using namespace std;

const int maxn = 1e6 + 100;
typedef long long ll;
int n, k, q;
int t[maxn], a[maxn];
ll pre[maxn];
int id[3][maxn];//id[0][i], id[1][i], id[2][i]分别表示pre[i], pre[i]+k, pre[i-k]离散化后的值

struct Query
{
    int id, l, r;
    ll ans;
} querys[maxn];

ll cnt[maxn];//cnt[i]中i是离散后的值
ll ans;
void update(int pos, int x, int add)
{
    if (add == 1)//区间增大
    {
        ans += cnt[x];//答案增加
        cnt[pos]++;//更新区间
    }
    else
    {
        cnt[pos]--;
        ans -= cnt[x];
    }
}

int block_size;
bool cmp(const Query&x, const Query&y)
{
    if ((x.l - 1) / block_size + 1 == (y.l - 1) / block_size + 1) return x.r < y.r;
    return x.l < y.l;
}
bool cmp_id(const Query&x, const Query&y)
{
    return x.id < y.id;
}

int main()
{
    scanf("%d%d", &n, &k);
    for (int i = 1; i <= n; ++i) scanf("%d", t + i);
    for (int i = 1; i <= n; ++i) scanf("%d", a + i);
    for (int i = 1; i <= n; ++i)
    {
        if (t[i] == 1) pre[i] = pre[i - 1] + a[i];
        else pre[i] = pre[i - 1] - a[i];
    }
    scanf("%d", &q);
    for (int i = 1; i <= q; ++i)
    {
        scanf("%d%d", &querys[i].l, &querys[i].r);
        querys[i].id = i;
    }
    //排序
    block_size = sqrt(n);
    sort(querys + 1, querys + 1 + q, cmp);
    //离散化
    int all = 0;
    map<ll, int> mp;
    mp[0] = all++, mp[k] = all++, mp[-k] = all++;
    for (int i = 0; i <= n; ++i)
    {
        mp[pre[i]] = all++;
        mp[pre[i] - k] = all++;
        mp[pre[i] + k] = all++;
    }
    for (int i = 0; i <= n; ++i)
    {
        id[0][i] = mp[pre[i]];
        id[1][i] = mp[pre[i] + k];
        id[2][i] = mp[pre[i] - k];
    }

    int l = 1, r = 0;
    cnt[mp[0]] = 1;
    ans = 0;
    for (int i = 1; i <= q; ++i)
    {
        for (; l < querys[i].l; ++l) update(id[0][l - 1], id[1][l - 1], -1);
        //区间从[l, r]变成[l+1, r], cnt统计的区间从[l-1, r]变成[l, r]
        //pre[l-1]删去, 原答案中包含pre[l-1]的都要减去, 也就是要减去满足pre[x] - pre[l-1] = k的x的个数
        //所以向update函数传入pre[x](pre[l-1]+k)和pre[l-1]的离散后的值, 更新答案, 其余三个循环与此类是

        for (; l > querys[i].l; --l) update(id[0][l - 2], id[1][l - 2], 1);
        for (; r < querys[i].r; ++r) update(id[0][r + 1], id[2][r + 1], 1);
        for (; r > querys[i].r; --r) update(id[0][r], id[2][r], -1);

        querys[i].ans = ans;
    }
    sort(querys + 1, querys + 1 + q, cmp_id);
    for (int i = 1; i <= q; ++i) printf("%lld\n", querys[i].ans);

    return 0;
}