【題目連結】

【前置技能】

• 字尾自動機/字尾樹

【題解】

• 發現葉節點的個數非常少，所以建立葉節點路徑上的字串的廣義字尾自動機，然後統計一下本質不同的子串的個數即可。
• 時間複雜度$O(N*X)$($X$為葉節點個數)

【程式碼】

#include<bits/stdc++.h>
#define INF 0x3f3f3f3f
#define LL  long long
#define MAXN    100010
using namespace std;
int n, c, w[MAXN], is_leaf[MAXN], leaf[30], cnt, sta[MAXN], top, cur;
LL ans;
vector <int> a[MAXN];
 
template <typename T> void chkmin(T &x, T y){x = min(x, y);}
template <typename T> void chkmax(T &x, T y){x = max(x, y);}
template <typename T> void read(T &x){
    x = 0; int f = 1; char ch = getchar();
    while (!isdigit(ch)) {if (ch == '-') f = -1; ch = getchar();}
    while (isdigit(ch)) {x = x * 10 + ch - '0'; ch = getchar();}
    x *= f;
}
 
struct Suffix_Automaton{
    struct info{int son[10], fa, dep;}a[MAXN * 50 * 2];
    int root, cnt, last;
    int create(int d){
        a[++cnt].dep = d;
        return cnt;
    }
    void extend(int ch){
        int p = last, q = a[p].son[ch];
        if (!q) {
            int np = create(a[p].dep + 1);
            while (!a[p].son[ch]){
                a[p].son[ch] = np;
                p = a[p].fa;
            }
            if (a[p].son[ch] == np) a[np].fa = root;
            else {
                q = a[p].son[ch];
                if (a[q].dep == a[p].dep + 1) a[np].fa = q;
                else {
                    int nq = create(a[p].dep + 1);
                    memcpy(a[nq].son, a[q].son, sizeof(a[q].son));
                    a[nq].fa = a[q].fa;
                    a[q].fa = a[np].fa = nq;
                    while (a[p].son[ch] == q){
                        a[p].son[ch] = nq;
                        p = a[p].fa;
                    }
                }
            }
            last = np;
        } else {
            if (a[q].dep == a[p].dep + 1) last = q;
            else {
                int nq = create(a[p].dep + 1);
                memcpy(a[nq].son, a[q].son, sizeof(a[q].son));
                a[nq].fa = a[q].fa;
                a[q].fa = nq;
                while (a[p].son[ch] == q){
                    a[p].son[ch] = nq;
                    p = a[p].fa;
                }
                last = nq;
            }
        }
    }
    void insert(int *s, int len){
        last = 0;
        for (int i = 1; i <= len; ++i)
            extend(s[i]);
    }
    void getans(){
        for (int i = 1; i <= cnt; ++i)
            ans += a[i].dep - a[a[i].fa].dep;
    }
}SAM;
 
void dfs(int pos, int dad){
    sta[++top] = w[pos];
    if (is_leaf[pos] && pos != cur) SAM.insert(sta, top);
    for (unsigned i = 0, si = a[pos].size(); i < si; ++i){
        int son = a[pos][i];
        if (son != dad){
            dfs(son, pos);
        }
    }
    --top;
}
 
int main(){
    read(n), read(c);
    for (int i = 1; i <= n; ++i)
        read(w[i]);
    for (int i = 1; i < n; ++i){
        int u, v; read(u), read(v);
        a[u].push_back(v);
        a[v].push_back(u);
    } 
    for (int i = 1; i <= n; ++i)
        if (a[i].size() == 1) is_leaf[i] = 1, leaf[++cnt] = i;
    for (int i = 1; i <= cnt; ++i){
        cur = leaf[i];
        dfs(leaf[i], 0);
    }
    SAM.getans();
    printf("%lld\n", ans);
    return 0;
}