【后缀自动机】 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;
}