P5161 WD与数列（后缀自动机+线段树合并）

Posted 2021-01-31 bztminamoto

传送门

没想出来→_→
首先不难看出要差分之后计算不相交也不相邻的相等子串对数，于是差分之后建SAM，在parent树上用线段树合并维护endpos集合，然后用启发式合并维护一个节点对另一个节点的贡献，于是总的时间复杂度为(O(nlog^2n))

//minamoto
#include<bits/stdc++.h>
#define R register
#define ll long long
#define IT vector<int>::iterator
#define fp(i,a,b) for(R int i=a,I=b+1;i<I;++i)
#define fd(i,a,b) for(R int i=a,I=b-1;i>I;--i)
#define go(u) for(int i=head[u],v=e[i].v;i;i=e[i].nx,v=e[i].v)
#define gg(u) for(IT it=f[u].begin();it!=f[u].end();++it)
using namespace std;
char buf[1<<21],*p1=buf,*p2=buf;
inline char getc(){return p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++;}
int read(){
    R int res,f=1;R char ch;
    while((ch=getc())>'9'||ch<'0')(ch=='-')&&(f=-1);
    for(res=ch-'0';(ch=getc())>='0'&&ch<='9';res=res*10+ch-'0');
    return res*f;
}
const int N=6e5+5;
struct eg{int v,nx;}e[N];int head[N],tot;
inline void add(R int u,R int v){e[++tot]={v,head[u]},head[u]=tot;}
int fa[N],l[N],a[N],rt[N];map<int,int>ch[N];
struct node{int s,ls,rs;ll xs;}t[N<<5];vector<int>*f[N],tmp[N];
int n,m,ctot,cnt,las;ll ans;
void ins(int c){
    int p=las,np=++cnt;las=np,l[np]=l[p]+1;
    for(;p&&!ch[p].count(c);p=fa[p])ch[p][c]=np;
    if(!p)fa[np]=1;
    else{
        int q=ch[p][c];
        if(l[q]==l[p]+1)fa[np]=q;
        else{
            int nq=++cnt;l[nq]=l[p]+1;
            ch[nq]=ch[q],fa[nq]=fa[q],fa[np]=fa[q]=nq;
            for(;p&&ch[p][c]==q;p=fa[p])ch[p][c]=nq;
        }
    }
}
void update(int &p,int l,int r,int x){
    if(!p)p=++ctot;++t[p].s,t[p].xs+=x;
    if(l==r)return;
    int mid=(l+r)>>1;
    x<=mid?update(t[p].ls,l,mid,x):update(t[p].rs,mid+1,r,x);
}
int qaqs(int p,int l,int r,int ql,int qr){
    if(!p)return 0;if(ql<=l&&qr>=r)return t[p].s;
    int mid=(l+r)>>1,res=0;
    if(ql<=mid)res+=qaqs(t[p].ls,l,mid,ql,qr);
    if(qr>mid)res+=qaqs(t[p].rs,mid+1,r,ql,qr);
    return res;
}
inline int querys(R int p,R int l,R int r){return (l>r||r<1||l>n)?0:qaqs(p,1,n,l,r);}
ll qaqxs(int p,int l,int r,int ql,int qr){
    if(!p)return 0;if(ql<=l&&qr>=r)return t[p].xs;
    int mid=(l+r)>>1;ll res=0;
    if(ql<=mid)res+=qaqxs(t[p].ls,l,mid,ql,qr);
    if(qr>mid)res+=qaqxs(t[p].rs,mid+1,r,ql,qr);
    return res;
}
inline ll queryxs(R int p,R int l,R int r){return (l>r||r<1||l>n)?0:qaqxs(p,1,n,l,r);}
int merge(int x,int y){
    if(!x||!y)return x|y;
    t[x].s+=t[y].s,t[x].xs+=t[y].xs;
    t[x].ls=merge(t[x].ls,t[y].ls),t[x].rs=merge(t[x].rs,t[y].rs);
    return x;
}
void dfs(int u){
    int k=l[u];
    go(u){
        dfs(v);if(f[u]->size()<f[v]->size())swap(f[u],f[v]),swap(rt[u],rt[v]);
        gg(v){
            int x=*it;
            ll A=1ll*querys(rt[u],x-k-1,x-2)*(x-1)-queryxs(rt[u],x-k-1,x-2)+1ll*querys(rt[u],1,x-k-2)*k;
            ll B=1ll*queryxs(rt[u],x+2,x+k+1)-1ll*querys(rt[u],x+2,x+k+1)*(x+1)+1ll*querys(rt[u],x+k+2,n)*k;
            ans+=A+B,f[u]->push_back(x);
        }rt[u]=merge(rt[u],rt[v]);
    }
}
int main(){
//  freopen("testdata.in","r",stdin);
    n=read();
    fp(i,1,n)a[i]=read();
    fp(i,1,n-1)a[i]=a[i+1]-a[i],f[i]=&tmp[i];
    ans=1ll*n*(n-1)>>1;
    --n,las=cnt=1;
    fp(i,1,n)ins(a[i]),f[las].push_back(i),update(rt[las],1,n,i);
    fp(i,2,cnt)add(fa[i],i);
    dfs(1);printf("%lld
",ans);
    return 0;
}