题解：CF2121F Yamakasi

· · 题解

\color{blue}{\texttt{[Problem Discription]}}

给定一个长度为 n 的序列 a_{1 \dots n}，求满足以下三个条件的区间 [l,r] 的数量：

多组数据，保证 n 的总和不超过 2 \times 10^{5}。

\color{blue}{\texttt{[Analysis]}}

首先，涉及连续子区间求和的问题很容易想到前缀和。

我们把数轴上 1,2,3, \dots, n 这 n 个点按前缀和分为若干组，把前缀和相同的归入一组，前缀和不同的点归为不同的组。这样，在区间左端点 l 确定的情况下，满足第一个条件的区间右端点在同一个集合里面。

不妨设 r 的候选集合为 S，我们将 S 中所有的元素从小到大排序，并形式化的记为：

x_{1}<x_{2}<\dots<x_{m} \quad \text{s.t.} \quad m\in \mathbb{N},x_{i} \in S

显然，最大值是有单调性的，即如果 x_{s}\geq l，则有

\max\limits_{i=l}^{x_{s}} \{ a_{i} \} \leq \max\limits_{i=l}^{x_{s+1}} \{ a_{i} \}

根据单调性，在候选区间中满足第二个条件的子集一定是原候选区间的一个连续的子区间，而且可以使用二分法把这个区间找出来（类似于 lower_bound 和 upper_bound 那样）。

用 ST 表求出区间最大值，总时间复杂度 O(n \log n)。

\color{blue}{\text{Code}} 
template<typename T>
T ckmax(T &a,T b){
    return (a=((a<b)?b:a));
}

const int N=2e5+100;

vector<int> pos[N];
int n,T,a[N],maxn;
ll sum,ans,pre[N];

ll block[N];
int blcnt,b[N];

class ST_table{
    public:
        void ST_init(int *a,int n){
            _log[0]=-1;
            for(int i=1;i<=n;i++)
                _log[i]=_log[i>>1]+1;

            for(int i=1;i<=n;i++)
                f[0][i]=a[i];

            for(int j=1;j<=_log[n];j++)
                for(int i=1;i+(1<<j)-1<=n;i++)
                    f[j][i]=max(f[j-1][i],f[j-1][i+(1<<(j-1))]);
        }
        int query(int l,int r){
            int s=_log[r-l+1];
            return max(f[s][l],f[s][r-(1<<s)+1]);
        }
    private:
        int f[22][N],_log[N];
};//维护一个 ST 表

ST_table st;

int Bineary_Search_low(int cur){
    int Pos=lower_bound(block+1,block+blcnt+1,pre[cur-1]+sum)-block;
    if (Pos>blcnt||block[Pos]!=pre[cur-1]+sum) return -1;//不存在这个前缀和

    int ans=-1,r=pos[Pos].size()-1;
    int l=lower_bound(pos[Pos].begin(),pos[Pos].end(),cur)-pos[Pos].begin(); 
    while (l<=r){
        int mid=(l+r)>>1;
        if (st.query(cur,pos[Pos][mid])>=maxn){
            ans=mid;
            r=mid-1;
        }
        else l=mid+1;
    }

    return ans;
}//在候选区间中二分求出可行域左端点
int Bineary_Search_high(int cur){
    int Pos=lower_bound(block+1,block+blcnt+1,pre[cur-1]+sum)-block;
    if (Pos>blcnt||block[Pos]!=pre[cur-1]+sum) return -1;

    int ans=-1,r=pos[Pos].size()-1;
    int l=lower_bound(pos[Pos].begin(),pos[Pos].end(),cur)-pos[Pos].begin(); 
    while (l<=r){
        int mid=(l+r)>>1;
        if (st.query(cur,pos[Pos][mid])<=maxn){
            ans=mid;
            l=mid+1;
        }
        else r=mid-1;
    }

    return ans;
}//在候选区间中二分求出可行域右端点

int OneZDoubleY(){
    n=read();
    sum=read();maxn=read();
    for(int i=1;i<=n;i++)
        a[i]=read();

    st.ST_init(a,n);

    ans=blcnt=0;

    for(int i=1;i<=n;i++){
        pre[i]=pre[i-1]+a[i];
        block[++blcnt]=pre[i];
    }

    sort(block+1,block+blcnt+1);
    blcnt=unique(block+1,block+blcnt+1)-block-1;
    for(int i=1;i<=n;i++)
        b[i]=lower_bound(block+1,block+blcnt+1,pre[i])-block;

    for(int i=1;i<=blcnt;i++)
        pos[i].clear();

    for(int i=1;i<=n;i++)
        pos[b[i]].push_back(i);//维护前缀和相同的点的集合

    for(int i=1;i<=n;i++){
        int l=Bineary_Search_low(i);
        int r=Bineary_Search_high(i);

        if (l!=-1&&r!=-1&&l<=r) ans+=r-l+1;
    }

    cout<<ans<<endl;
    return 0;
}

int main(){
    T=read();
    while (T--) OneZDoubleY();

    return 0;
}

有这么一个坑点，可能会导致 TLE。

int Bineary_Search_low(int cur){
    int Pos=lower_bound(block+1,block+blcnt+1,pre[cur-1]+sum)-block;
    if (Pos>blcnt||block[Pos]!=pre[cur-1]+sum) return -1;

    vector<int> tmp=pos[Pos];//(*)

    int ans=-1,r=tmp.size()-1;
    int l=lower_bound(tmp.begin(),tmp.end(),cur)-tmp.begin(); 
    while (l<=r){
        int mid=(l+r)>>1;
        if (st.query(cur,tmp[mid])>=maxn){
            ans=mid;
            r=mid-1;
        }
        else l=mid+1;
    }

    return ans;
}

在星号标记处，为了后面书写的方便，我把 pos 的值传递给了 tmp。但是这份代码是会 TLE 的。

为什么呢？

这就要涉及到 vector 的底层实现问题了。vector 间相互赋值的时候相当于是要把所有的值拷贝一份，是 O(N) 的，其中 N 是 vector 的长度。

在极端情况下，如 a_{1 \dots n} 中所有元素均为 0，vector 的长度为 O(n)，赋值进行了 O(n) 次，因此总的时间复杂度为 O(n^{2})，当然会 TLE。