Candies 【后缀数组】【主席树】

题目链接：https://vjudge.net/problem/Gym-102302K

前置技能：后缀数组处理本质不同的字串。

#include <bits/stdc++.h>
#define rep(i, a, b) for(int i = (a); i <= (b); i++)
#define per(i, a, b) for(int i = (a); i >= (b); i--)
#define pb push_back
#define mp make_pair
#define mid (l+r>>1)
#define lson (o<<1)
#define rson (o<<1|1)
#define all(x) (x).begin(),(x).end()
#define ll long long
using namespace std;
const int MAX = 3e6+1000;
const int N = 3e6+1000;
const int SIGMA_SIZE = 3e6+1000;
int cntA[MAX], cntB[MAX], sa[MAX], tsa[MAX], A[MAX], B[MAX], height[MAX], Rank[MAX];
int n, data[MAX];
int bucket[SIGMA_SIZE];

int s[MAX];
void get_SA(int *ch) {
    for (int i = 0; i < SIGMA_SIZE; i++) bucket[i] = 0;
    for (int i = 1; i <= n; i++) bucket[ch[i - 1]]++;
    for (int i = 1; i < SIGMA_SIZE; i++) bucket[i] += bucket[i - 1];
    for (int i = n; i; i--) sa[bucket[ch[i - 1]]--] = i;
    Rank[sa[1]] = 1;
    for (int i = 2; i <= n; i++) {
        Rank[sa[i]] = Rank[sa[i - 1]];
        if (ch[sa[i] - 1] != ch[sa[i - 1] - 1]) Rank[sa[i]]++;
    }
    for (int l = 1; Rank[sa[n]] < n; l <<= 1) {
        memset(cntA, 0, sizeof(cntA));
        memset(cntB, 0, sizeof(cntB));
        for (int i = 1; i <= n; i++) {
            cntA[A[i] = Rank[i]]++;
            cntB[B[i] = (i + l <= n) ? Rank[i + l] : 0]++;
        }
        for (int i = 1; i <= n; i++) cntB[i] += cntB[i - 1];
        for (int i = n; i; i--) tsa[cntB[B[i]]--] = i;
        for (int i = 1; i <= n; i++) cntA[i] += cntA[i - 1];
        for (int i = n; i; i--) sa[cntA[A[tsa[i]]]--] = tsa[i];
        Rank[sa[1]] = 1;
        for (int i = 2; i <= n; i++) {
            Rank[sa[i]] = Rank[sa[i - 1]];
            if (A[sa[i]] != A[sa[i - 1]] || B[sa[i]] != B[sa[i - 1]]) Rank[sa[i]]++;
        }
    }
    for (int i = 1, j = 0; i <= n; i++) {
        if (j) j--;
        while (ch[i + j - 1] == ch[sa[Rank[i] - 1] + j - 1]) j++;
        height[Rank[i]] = j;
    }
}
//--------------------------------
struct TREE{
    int l,r,x;
}tree[N*20];
int tot,rt[N];
void insert(int &o,int l,int r,int data){
    tree[++tot] = tree[o];
    o = tot;
    tree[o].x++;
    if(l==r) return;
    if(mid>=data) insert(tree[o].l,l,mid,data);
    else insert(tree[o].r,mid+1,r,data);
}
ll query(int lo,int ro,int l,int r,int h,int t) {
    if(l>=h&&r<=t) return tree[ro].x-tree[lo].x;
    ll ans=0;
    if(mid>=h) ans+=query(tree[lo].l,tree[ro].l,l,mid,h,t);
    if(mid<t)  ans+=query(tree[lo].r,tree[ro].r,mid+1,r,h,t);
    return ans;
}

vector<ll> k;
ll L,R;
ll sum[N];


ll solve(int x,int ql,int qr) {
    int tl = lower_bound(all(k),L+sum[x-1])-k.begin()+1;
    int tr = lower_bound(all(k),R+sum[x-1])-k.begin()+1;
    ll t = query(rt[ql-1],rt[qr],1,N-100,tl,tr);
    return t;
}
int main() {
    //freopen("a.txt","r",stdin);
    ios::sync_with_stdio(0);
    cin>>n>>L>>R;
    rep(i, 0, n-1) {
        cin>>s[i];
        sum[i+1] = s[i]+sum[i];
        k.pb(s[i]);
        k.pb(sum[i+1]);
        k.pb(L+sum[i+1]);
        k.pb(R+sum[i+1]);
    }
    sort(all(k));
    k.erase(unique(all(k)),k.end());
    rep(i, 0, n-1) s[i] = lower_bound(all(k),s[i])-k.begin()+1;
    get_SA(s);

    rt[0] = 0;
    tree[0].l = tree[0].r = tree[0].x = 0;
    for(int i = 1;i <= n;i ++){
        rt[i] = rt[i-1];
        int x = lower_bound(all(k),sum[i])-k.begin()+1;
        insert(rt[i],1,N-100,x);
    }
    ll ans = 0;
    rep(i, 1, n) 
        ans += solve(sa[i],sa[i]+height[i],n);
    cout<<ans;
    return 0;
}