【BZOJ3926】【ZJOI2015】诸神眷顾的幻想乡（广义后缀自动机）

Description

幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的2600岁生日，无数幽香的粉丝到了幽香家门前的太阳花田上来为幽香庆祝生日。
粉丝们非常热情，自发组织表演了一系列节目给幽香看。幽香当然也非常高兴啦。
这时幽香发现了一件非常有趣的事情，太阳花田有n块空地。在过去，幽香为了方便，在这n块空地之间修建了n-1条边将它们连通起来。也就是说，这n块空地形成了一个树的结构。
有n个粉丝们来到了太阳花田上。为了表达对幽香生日的祝贺，他们选择了c中颜色的衣服，每种颜色恰好可以用一个0到c-1之间的整数来表示。并且每个人都站在一个空地上，每个空地上也只有一个人。这样整个太阳花田就花花绿绿了。幽香看到了，感觉也非常开心。
粉丝们策划的一个节目是这样的，选中两个粉丝A和B（A和B可以相同），然后A所在的空地到B所在的空地的路径上的粉丝依次跳起来（包括端点），幽香就能看到一个长度为A到B之间路径上的所有粉丝的数目（包括A和B）的颜色序列。一开始大家打算让人一两个粉丝（注意：A,B和B,A是不同的，他们形成的序列刚好相反，比如红绿蓝和蓝绿红）都来一次，但是有人指出这样可能会出现一些一模一样的颜色序列，会导致审美疲劳。
于是他们想要问题，在这个树上，一共有多少可能的不同的颜色序列（子串）幽香可以看到呢？
太阳花田的结构比较特殊，只与一个空地相邻的空地数量不超过20个。

---

Solution

发现叶节点非常少，从每个叶节点出发将得到的串插入广义后缀自动机，然后记数就行了QAQ

---

Code

// luogu-judger-enable-o2
/************************************************
 * Au: Hany01
 * Date: May 8th, 2018
 * Prob: [BZOJ3926][ZJOI2015] 诸神眷顾的幻想乡
 * Email: [email protected]
************************************************/

#include<bits/stdc++.h>

using namespace std;

typedef long long LL;
typedef pair<int, int> PII;
#define File(a) freopen(a".in", "r", stdin), freopen(a".out", "w", stdout)
#define rep(i, j) for (register int i = 0, i##_end_ = (j); i < i##_end_; ++ i)
#define For(i, j, k) for (register int i = (j), i##_end_ = (k); i <= i##_end_; ++ i)
#define Fordown(i, j, k) for (register int i = (j), i##_end_ = (k); i >= i##_end_; -- i)
#define Set(a, b) memset(a, b, sizeof(a))
#define Cpy(a, b) memcpy(a, b, sizeof(a))
#define x first
#define y second
#define pb(a) push_back(a)
#define mp(a, b) make_pair(a, b)
#define ALL(a) (a).begin(), (a).end()
#define SZ(a) ((int)(a).size())
#define INF (0x3f3f3f3f)
#define INF1 (2139062143)
#define Mod (1000000007)
#define debug(...) fprintf(stderr, __VA_ARGS__)
#define y1 wozenmezhemecaia

template <typename T> inline bool chkmax(T &a, T b) { return a < b ? a = b, 1 : 0; }
template <typename T> inline bool chkmin(T &a, T b) { return b < a ? a = b, 1 : 0; }

inline int read()
{
    register int _, __; register char c_;
    for (_ = 0, __ = 1, c_ = getchar(); c_ < '0' || c_ > '9'; c_ = getchar()) if (c_ == '-') __ = -1;
    for ( ; c_ >= '0' && c_ <= '9'; c_ = getchar()) _ = (_ << 1) + (_ << 3) + (c_ ^ 48);
    return _ * __;
}

const int maxn = 100005, maxc = 2000005;

int n, c, ch[maxc][10], v[maxn << 1], nex[maxn << 1], beg[maxn], e, deg[maxn], len[maxc], fa[maxc], co[maxn], tot = 1, lst[maxn], cnt;

inline void add(int uu, int vv) { v[++ e] = vv, nex[e] = beg[uu], beg[uu] = e, ++ deg[vv]; }

inline int extend(int c, int p)
{
    int np = ++ tot;
    len[np] = len[p] + 1;
    while (p && !ch[p][c]) ch[p][c] = np, p = fa[p];
    if (!p) fa[np] = 1;
    else {
        int q = ch[p][c];
        if (len[q] == len[p] + 1) fa[np] = q;
        else {
            int nq = ++ tot;
            fa[nq] = fa[q], Cpy(ch[nq], ch[q]), len[nq] = len[p] + 1;
            fa[q] = fa[np] = nq;
            while (p && ch[p][c] == q) ch[p][c] = nq, p = fa[p];
        }
    }
    return np;
}

void dfs(int u, int fa, int p)
{
    p = extend(co[u], p);
    if (!nex[beg[u]] && v[beg[u]] == fa) return ;
    for (register int i = beg[u]; i; i = nex[i])
        if (v[i] != fa) dfs(v[i], u, p);
}

int main()
{
#ifdef hany01
    File("bzoj3926");
#endif

    static int uu, vv;

    n = read(), c = read();
    For(i, 1, n) co[i] = read();
    For(i, 2, n)
        uu = read(), vv = read(), add(uu, vv), add(vv, uu);

    For(i, 1, n)
        if (deg[i] == 1) dfs(i, 0, 1);

    static LL Ans = 0;
    For(i, 1, tot) Ans += len[i] - len[fa[i]];
    printf("%lld\n", Ans);

    return 0;
}
//西风稍急喧窗竹，停又续，腻脸悬双玉。
//    -- 阎选《河传·秋雨》