【HAOI2016】找相同字符

#include<cstdio>
#include<cstring>
#include<iostream>
#include<vector>
#define N 200050
#define ri register int
using namespace std;
char s1[N<<1],s2[N<<1];
int l1,l2;
int cnt[N<<1];
vector<int> son[N<<1];
long long ans=0;

struct SAM{
  int ff[N<<1],len[N<<1];
  int ch[N<<1][26];
  int par[N<<1][20];
  int ton[N<<1];
  int las,tot;
  void init() {
    las=tot=1;
  }
  void extend(int c){
    int np=++tot,p=las; las=np;
    len[np]=len[p]+1;
    cnt[np]++;
    while (p && !ch[p][c]) ch[p][c]=np,p=ff[p];
    if (!p) ff[np]=1;
    else {
      int q=ch[p][c];
      if (len[q]==len[p]+1) ff[np]=q;
      else {
        int nq=++tot;
        len[nq]=len[p]+1;
        for (ri i=0;i<26;i++) ch[nq][i]=ch[q][i]; ff[nq]=ff[q];
        ff[np]=ff[q]=nq;
        while (p && ch[p][c]==q) ch[p][c]=nq,p=ff[p];
      }
    }
  }
  void maketree() {
    for (ri i=1;i<=tot;i++) son[ff[i]].push_back(i);
  }
  void init_par(int x) {
    par[x][0]=ff[x];
    for (ri i=1;i<=19;i++) par[x][i]=par[par[x][i-1]][i-1];
    for (ri i=0;i<son[x].size();i++) {
      init_par(son[x][i]);
      cnt[x]+=cnt[son[x][i]];
    }
  }
  
  void match(char *s){
    int n=strlen(s+1);
    int now=1;
    int cc=0;
    for (ri i=1;i<=n;i++) {
      if (ch[now][s[i]-‘a‘]) now=ch[now][s[i]-‘a‘],cc++;
      else {
        while (!ch[now][s[i]-‘a‘]) now=ff[now];
        if (!now) {
          now=1; cc=0;
        }
        else {
          cc=len[now]+1;
          now=ch[now][s[i]-‘a‘];
        }
      }
      int x=now;
      ans+=(cc-len[ff[x]]*1LL)*1LL*cnt[x];
      ton[ff[x]]++;
    }
  }
  
  void dp(int x){
    for (ri i=0;i<son[x].size();i++) {
      dp(son[x][i]);
      ton[x]+=ton[son[x][i]];
    }
    ans+=cnt[x]*1LL*ton[x]*1LL*(len[x]-len[ff[x]]);
  }
} sam;

int main() {
  scanf("%s%s",s1+1,s2+1);
  l1=strlen(s1+1);
  sam.init();
  for (ri i=1;i<=l1;i++) sam.extend(s1[i]-‘a‘);
  sam.maketree();
  sam.init_par(1);
  sam.match(s2);
  sam.dp(1);
  printf("%lld\n",ans);
}