洛谷P3346 [ZJOI2015]诸神眷顾的幻想乡(广义后缀自动机)

#include<bits/stdc++.h>
#define LL long long 
using namespace std;
const int MAXN = 2e6 + 10;
inline int read() {
    char c = getchar(); int x = 0, f = 1;
    while(c < '0' || c > '9') {if(c == '-') f = -1; c = getchar();}
    while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
    return x * f;
}
int N, C, a[MAXN], deg[MAXN];
namespace SAM {
    int fa[MAXN], len[MAXN], ch[MAXN][11], tot = 1, root = 1, las = 1;
    int insert(int x, int pre) {
        int now = ++tot; len[now] = len[pre] + 1;
        for(; pre && !ch[pre][x]; pre = fa[pre]) ch[pre][x] = now;
        if(!pre) {fa[now] = root; return now;}
        int q = ch[pre][x];
        if(len[q] == len[pre] + 1) fa[now] = q;
        else {
            int nq = ++tot; fa[nq] = fa[q]; len[nq] = len[pre] + 1;
            memcpy(ch[nq], ch[q], sizeof(ch[q]));
            for(; pre && ch[pre][x] == q; pre = fa[pre]) ch[pre][x] = nq;
            fa[q] = fa[now] = nq;
        }
        return now;
    }
    LL calc() {
        LL ans = 0;
        for(int i = 1; i <= tot; i++) ans += len[i] - len[fa[i]];
        return ans;
    }
}
vector<int> v[MAXN], node;
void dfs(int x, int _fa, int p) {
    p = SAM::insert(a[x], p);
    for(auto &to : v[x]) {
        if(to == _fa) continue;
        dfs(to, x, p);
    }
}   
int main() {
    N = read(); C = read();
    for(int i = 1; i <= N; i++) a[i] = read();
    for(int i = 1; i <= N - 1; i++) {
        int x = read(), y = read();
        v[x].push_back(y); deg[x]++;
        v[y].push_back(x); deg[y]++;
    }
    for(int i = 1; i <= N; i++) 
        if(deg[i] == 1) SAM::las = 1, dfs(i, 0, 1);
    cout << SAM::calc();
    return 0;
}