省选模版复习——后缀数组

bzoj3238 后缀数组水

#include <cstdio>
#include <cstring>
#include <algorithm>
 
using namespace std;
const int Maxn=500005;
typedef long long LL;
LL Ans;
char S[Maxn];
int H[Maxn],Rank[Maxn],q[Maxn],pre[Maxn],nxt[Maxn];
int l,r,N,i,j,k;
struct Suf_Arr
{
  int x,num;
  bool operator <(const Suf_Arr &a)const
    { return x<a.x; }
} Sa[Maxn];
 
void calc_height(){
  for (i=1,k=0;i<=N;H[Rank[i++]]=k)
    for (k=max(k-1,0),j=Sa[Rank[i]-1].num;S[i+k]==S[j+k];k++);
}
 
void Suffix_Array(){
  scanf("%s",S+1);
  N=strlen(S+1);
  for (i=1;i<=N;i++)
    Sa[i].num=i,Sa[i].x=S[i],Rank[i]=1;
  for (int E=1;E<=N*2;E<<=1){
    for (i=1,j=2;j<=N+1;j++)
    if (Rank[Sa[j].num]!=Rank[Sa[i].num]){
      sort(Sa+i,Sa+j);
      Rank[Sa[i].num]=Rank[Sa[i-1].num]+1;
      for (k=i+1;k<j;k++)
      if (Sa[k].x==Sa[k-1].x) Rank[Sa[k].num]=Rank[Sa[k-1].num];
        else Rank[Sa[k].num]=Rank[Sa[k-1].num]+1;
      i = j;
    }
     
    for (i=1;i<=N;i++)
    if (Sa[i].num+E>N) Sa[i].x=0;
      else Sa[i].x=Rank[Sa[i].num+E];
  }
  calc_height();
}
 
void calc(){
  H[0] = -1; H[N+1] = -1;
  q[l=r=1]=0;
  for (i=1;i<=N;i++){
    while (l<=r && H[q[r]]>H[i]) r--;
    pre[i] = q[r]; q[++r] = i;
  }
  q[l=r=1]=N+1;
  for (i=N;i>0;i--){
    while (l<=r && H[q[r]]>H[i]) r--;
    nxt[i] = q[r]; q[++r] = i;
  }
   
  for (i=1;i<=N;i++){
    Ans = Ans + (LL)(i-pre[i])*(nxt[i]-i)*H[i];
    if (H[i]==H[nxt[i]]) pre[nxt[i]] = pre[i];
  }
  Ans = (LL)N*(N+1)/2*(N-1) - 2*Ans;
  printf("%lld\n",Ans);
}
 
int main(){
  //freopen("3238.in","r",stdin);
  //freopen("3238.out","w",stdout);
  Suffix_Array();
  calc();
  return 0;
}