AT_abc272_f [ABC272F] Two Strings 题解

后缀数组基本应用，关键点在字典序上。

首先肯定要把两个字符串在末尾复制一遍，这样就包含所有 \(f(S,i)\) 和 \(f(T,j)\) 了。由于总共只有两个字符串，很难往 \(Trie\) 树上考虑，那就很自然地往 \(SA\) 上靠。发现直接把俩字符串拼到一起然后跑一遍 \(SA\) 就能得到所有 \(f\) 字典序排名了（虽然还带着一长串的后缀，但是字典序贪心的比较方式使得多余的后缀完全不影响，唯一的变化就是判断相等从 \(s=t\) 改为 \(lcp(s,t)>=n\) ）。所以直接对着处理好的 \(rk\) 数组从前往后遍历统计答案即可，注意相等的字符串要放在一起一块处理，因为答案要把相等的也算进去。

复杂度是后缀数组的 \(O(nlogn)\) 。

代码：

#include<iostream>
#include<vector>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
const int N=8e5+5;
string s,t;
int n;
int m,p,rk[N<<1],oldrk[N],sa[N<<1],id[N],cnt[N];
int height[N];
void SA(){
    int m=128;
    for(int i=1;i<=n;i++) cnt[rk[i]=s[i]]++;
    for(int i=1;i<=m;i++) cnt[i]+=cnt[i-1];
    for(int i=n;i>=1;i--) sa[cnt[rk[i]]--]=i;

    for(int w=1;;w<<=1,m=p){
        int cur=0;
        for(int i=n-w+1;i<=n;i++) id[++cur]=i;
        for(int i=1;i<=n;i++)
            if(sa[i]>w) id[++cur]=sa[i]-w;

        memset(cnt,0,sizeof cnt);
        for(int i=1;i<=n;i++) cnt[rk[i]]++;
        for(int i=1;i<=m;i++) cnt[i] += cnt[i-1];
        for(int i=n;i>=1;i--) sa[cnt[rk[id[i]]]--]=id[i];
        
        p=0;
        memcpy(oldrk,rk,sizeof(oldrk));
        for(int i=1;i<=n;i++){
            if(oldrk[sa[i]]==oldrk[sa[i-1]]&&oldrk[sa[i]+w]==oldrk[sa[i-1]+w]) rk[sa[i]]=p;
            else rk[sa[i]]=++p;
        }

        if(p==n) break;
    }

    for(int i=1,k=0;i<=n;i++){
        if(rk[i]==0) continue;
        if(k) k--;
        while(s[i+k]==s[sa[rk[i]-1]+k]) k++;
        height[rk[i]]=k;
    }
}
int mn[N][21];
inline int lcp(int l,int r){
    if(l==r) return n;
    l++;
    int len=__lg(r-l+1);
    return min(mn[l][len],mn[r-(1<<len)+1][len]);
}
int main(){
    ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
    cin>>n;
    cin>>s>>t;
    int sz=s.size();
    s=s+s;
    s.erase(sz*2-1,1);
    t=t+t;
    t.erase(sz*2-1,1);
    n=s.size()+t.size()+1;
    s=' '+s+'#'+t;
    SA();
    for(int i=1;i<=n;i++) mn[i][0]=height[i];
    for(int i=1;i<=20;i++){
        for(int j=1;j+(1<<i)-1<=n;j++){
            mn[j][i]=min(mn[j][i-1],mn[j+(1<<(i-1))][i-1]);
        }
    }
    ll ans=0;
    int cnt=0;
    for(int i=1,j=0;i<=n;i=j+1){
        int cnt1=0,cnt2=0;
        while(j<n&&lcp(i,j+1)>=sz){
            j++;
            if(sa[j]<=sz) cnt2++;
            else if(sa[j]>sz*2&&sa[j]<=sz*3) cnt1++;
        }
        cnt+=cnt2;
        ans+=(ll)cnt1*cnt;
    }
    cout<<ans;
    return 0;
}