bzoj 4566 找相同字符【后缀自动机】

好不容易喜欢和大家一起吃饭TAT……难过……
题目大意：从A、B中各取出一个子串，求有多少种方案，使得两子串相同（不同位置的串算多次）

对A建后缀自动机，用B去匹配，若一个点被匹配到了，则他所有祖先全都能走到，每个点的贡献是|S| * (mx[i] - mn[i] + 1)

#include <iostream>
#include <cstdio>
#include <cstring>
#define N 400005
#define LL long long
using namespace std;

int n,m,tot = 1,last = 1,p,np,q,nq;
int mx[N],par[N],son[N][26],P[N],sm[N];
LL S[N],f[N],ans;
char a[N],b[N];

int new_node(int x)
{
    mx[++ tot] = x;
    return tot;
}

void add(int x)
{
    p = last;
    np = new_node(mx[p] + 1);
    for (;p && !son[p][x];p = par[p]) son[p][x] = np;
    if (!p) par[np] = 1;
    else
    {
        q = son[p][x];
        if (mx[q] == mx[p] + 1) par[np] = q;
        else
        {
            nq = new_node(mx[p] + 1);
            memcpy(son[nq],son[q],sizeof(son[nq]));
            par[nq] = par[q],par[q] = par[np] = nq;
            for (;son[p][x] == q;p = par[p]) son[p][x] = nq;
        }
    }
    S[last = np] = 1;
}

int main()
{
    scanf("%s%s",a + 1,b + 1);
    n = strlen(a + 1),m = strlen(b + 1);
    for (int i = 1;i <= n;i ++) add(a[i] - 'a');
    for (int i = 1;i <= tot;i ++) sm[mx[i]] ++;
    for (int i = 1;i <= n;i ++) sm[i] += sm[i - 1];
    for (int i = 1;i <= tot;i ++) P[sm[mx[i]] --] = i;
    for (int i = tot;i > 1;i --) S[par[P[i]]] += S[P[i]];
    for (int i = 2,x;i <= tot;i ++) x = P[i],f[x] = f[par[x]] + S[x] * (mx[x] - mx[par[x]]);
    for (int i = 1,p = 1,l = 0;i <= m;i ++)
    {
        int x = b[i] - 'a';
        if (son[p][x]) p = son[p][x],l ++;
        else
        {
            while (p && !son[p][x]) p = par[p];
            if (!p) p = 1,l = 0;
            else l = mx[p] + 1,p = son[p][x];
        }
        ans += f[par[p]] + S[p] * (l - mx[par[p]]);
    }
    cout << ans << endl;

    return 0;
}