Bzoj3926: [Zjoi2015]诸神眷顾的幻想乡

Posted 2020-11-01 Cyhlnj

题面

传送门

Sol

求多个串的不同子串的个数

广义后缀自动机
也就是可以表示所有的串的所有的后缀的自动机

那么每次建一个串后，另外一个串接在初始节点下面建就好了

叶子节点最多\(20\)个
那么对于每个叶子结点遍历一遍树建立\(sam\)
注意回溯时要把\(last\)指回来

# include <bits/stdc++.h>
# define RG register
# define IL inline
# define Fill(a, b) memset(a, b, sizeof(a))
using namespace std;
typedef long long ll;
 
IL int Input(){
    RG int x = 0, z = 1; RG char c = getchar();
    for(; c < '0' || c > '9'; c = getchar()) z = c == '-' ? -1 : 1;
    for(; c >= '0' && c <= '9'; c = getchar()) x = (x << 1) + (x << 3) + (c ^ 48);
    return x * z;
}
 
const int maxn(1e5 + 5);
const int maxm(4e6 + 5);

int trans[10][maxm], fa[maxm], tot = 1, last, len[maxm];
int n, first[maxn], cnt, col[maxn], degree[maxn];
ll ans;

struct Edge{
    int to, next;
} edge[maxn << 1];

IL void Add(RG int u, RG int v){
    edge[cnt] = (Edge){v, first[u]}, first[u] = cnt++;
}

IL void Extend(RG int c, RG int p){
    RG int np = ++tot;
    len[last = np] = len[p] + 1;
    while(p && !trans[c][p]) trans[c][p] = np, p = fa[p];
    if(!p) fa[np] = 1;
    else{
        RG int q = trans[c][p];
        if(len[q] == len[p] + 1) fa[np] = q;
        else{
            RG int nq = ++tot;
            fa[nq] = fa[q], len[nq] = len[p] + 1;
            for(RG int i = 0; i < 10; ++i) trans[i][nq] = trans[i][q];
            fa[q] = fa[np] = nq;
            while(p && trans[c][p] == q) trans[c][p] = nq, p = fa[p];
        }
    }
}

IL void Dfs(RG int u, RG int ff, RG int p){
    Extend(col[u], p), p = last;
    for(RG int e = first[u]; e != -1; e = edge[e].next)
        if(edge[e].to != ff) Dfs(edge[e].to, u, p);
}

int main(RG int argc, RG char *argv[]){
    n = Input(), Input();
    for(RG int i = 1; i <= n; ++i) first[i] = -1, col[i] = Input();
    for(RG int i = 1; i < n; ++i){
        RG int u = Input(), v = Input();
        Add(u, v), Add(v, u), ++degree[u], ++degree[v];
    }
    for(RG int i = 1; i <= n; ++i)
        if(degree[i] == 1) last = 1, Dfs(i, 0, 1);
    for(RG int i = 1; i <= tot; ++i) ans += len[i] - len[fa[i]];
    printf("%lld\n", ans);
    return 0;
}