【后缀自动机】 BZOJ 3926 ZJOI2015 诸神眷顾的幻想乡

好神的题。。。。

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
 
const int alpha = 26;
const int maxn = 4000005;
const int maxm = 200005;
 
struct node
{
    int len;
    node *ch[alpha], *fa;
}pool[maxn], *last[maxn], *root;
 
struct Edge
{
    int v;
    Edge *next;
}E[maxm << 1], *H[maxn], *edges;
 
int tot, n;
int a[maxm];
int du[maxm];
 
node *newnode(int len)
{
    node *p = &pool[tot++];
    p->len = len;
    p->fa = 0;
    memset(p->ch, 0, sizeof p->ch);
    return p;
}
 
void addedges(int u, int v)
{
    edges->v = v;
    edges->next = H[u];
    H[u] = edges++;
}
 
void init()
{
    tot = 0;
    last[0] = root = newnode(0);
 
    edges = E;
    memset(H, 0, sizeof H);
    memset(du, 0, sizeof du);
}
 
node* add(node *p, int c)
{
    node *np = p->ch[c];
    if(np && p->len + 1 == np->len) return np;
    else np = newnode(p->len + 1);
 
    for(; p && !p->ch[c]; p = p->fa) p->ch[c] = np;
    if(!p) np->fa = root;
    else {
        node *q = p->ch[c];
        if(q->len == p->len + 1) np->fa = q;
        else {
            node *nq = newnode(p->len + 1);
            *nq = *q;
            nq->len = p->len + 1;
            np->fa = q->fa = nq;
            for(; p && p->ch[c] == q; p = p->fa) p->ch[c] = nq;
        }
    }
    return np;
}
 
void dfs(int u, int fa)
{
    last[u] = add(last[fa], a[u]);
    for(Edge *e = H[u]; e; e = e->next) if(e->v != fa) dfs(e->v, u);
}
 
void work()
{
    int u, v;
    for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
     
    for(int i = 1; i < n; i++) {
        scanf("%d%d", &u, &v);
        addedges(u, v);
        addedges(v, u);
        du[u]++, du[v]++;
    }
     
    for(int i = 1; i <= n; i++) if(du[i] == 1) dfs(i, 0);
    LL ans = 0;
    for(int i = 1; i < tot; i++) ans += pool[i].len - pool[i].fa->len;
    printf("%lld\n", ans);
}
 
int main()
{
    while(scanf("%d%*d", &n) != EOF) {
        init();
        work();
    }
     
    return 0;
}