BZOJ 3926: [Zjoi2015]诸神眷顾的幻想乡(广义后缀自动机 多串)

因为任何一条路径都可以看做某两个叶子节点之间路径的一部分,然后分别把20个叶节点当作根,把整棵树看作trie树,那么一条路径就能看作是从根到某个点这一条路的后缀,构建SAM就能维护不同子串的个数了.

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
char cb[1<<15],*cs=cb,*ct=cb;
#define getc() (cs==ct&&(ct=(cs=cb)+fread(cb,1,1<<15,stdin),cs==ct)?0:*cs++)
template<class T>inline void read(T &res) {
    char ch; int flg = 1; while(!isdigit(ch=getc()))if(ch=='-')flg=-flg;
    for(res=ch-'0';isdigit(ch=getc());res=res*10+ch-'0'); res*=flg;
}
const int MAXN = 100005;
const int MAXM = 2000005<<1;
const int MAXC = 10;
int n, C, col[MAXN], deg[MAXN], fir[MAXN];
int to[MAXN<<1], nxt[MAXN<<1], cnt;
inline void add(int u, int v) { to[++cnt] = v, nxt[cnt] = fir[u], fir[u] = cnt; }
namespace Suffix_AutoMaton {
    int sz, last, len[MAXM], link[MAXM], ch[MAXM][MAXC];
    inline void init() {
        sz = last = 0; ++sz;
        len[0] = 0; link[0] = -1;
        //memset(ch, 0, sizeof ch);
    }
    inline void Copy(int A, int B) {
        link[A] = link[B];
        memcpy(ch[A], ch[B], sizeof ch[B]);
    }
    inline int insert(int p, int c) {
        if(ch[p][c]) {
            int q = ch[p][c];
            if(len[p] + 1 == len[q]) last = q;
            else {
                int x = sz++; last = x;
                len[x] = len[p] + 1;
                Copy(x, q);
                link[q] = x;
                for(; ~p && ch[p][c] == q; p = link[p]) ch[p][c] = x;
            }
        }
        else {
            int cur = sz++; last = cur;
            len[cur] = len[p] + 1;
            for(; ~p && !ch[p][c]; p = link[p]) ch[p][c] = cur;
            if(p == -1) link[cur] = 0;
            else {
                int q = ch[p][c];
                if(len[p] + 1 == len[q]) link[cur] = q;
                else {
                    int x = sz++;
                    len[x] = len[p] + 1;
                    Copy(x, q);
                    link[cur] = link[q] = x;
                    for(; ~p && ch[p][c] == q; p = link[p]) ch[p][c] = x;
                }
            }
        }
        return last;
    }
    void dfs(int x, int ff, int p) {
        p = insert(p, col[x]);
        for(int i = fir[x]; i; i = nxt[i])
            if(to[i] != ff) dfs(to[i], x, p);
    }
    inline LL getans() {
        LL re = 0;
        for(int i = 1; i < sz; ++i)
            re += len[i] - len[link[i]];
        return re;
    }
}

int main() {
    read(n), read(C);
    for(int i = 1; i <= n; ++i) read(col[i]);
    for(int i = 1, x, y; i < n; ++i) {
        read(x), read(y);
        ++deg[x], ++deg[y];
        add(x, y), add(y, x);
    }
    Suffix_AutoMaton::init();
    for(int i = 1; i <= n; ++i)
        if(deg[i] == 1)
            Suffix_AutoMaton::dfs(i, 0, 0);
    printf("%lld\n", Suffix_AutoMaton::getans());
}