P3346 [ZJOI2015] 诸神眷顾的幻想乡

[ZJOI2015] 诸神眷顾的幻想乡

题目描述

幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的 \(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\) 个。

数据规模与约定

• 对于 \(100\%\) 的数据，\(1\le n\le10^5\)，\(1\le c\le10\)，\(1 \leq u, v \leq n\)。

Soution:

感觉最近学的算法都好神,LCT一天四发实在写不动了，所以跑来水报告^-。

好神的 SAM 题。但是首先我们明确一个树上结论：对于一颗无根树，任意两点间的一个路径会在以某个叶子节点为根遍历时变为一条从根出发到叶子的路径的子路径。形式化的：设叶子集为 \(Lev\) , \(G_x\) 表示以 \(x\) 为根的遍历方式

\[\exists rt \in Lev,(x,y) \in G_{rt} ,(x,y) \subseteq (rt,l),l\in Lev ,l \ne rt \]

写的有点奇怪，但是意思时对的。

然后我们不难发现： 由于太阳花田的结构比较特殊，只与一个空地相邻的空地数量不超过 \(20\) 个。 就提示了我们枚举以 rt 来求出这棵树上的所有路径，我们用这样的方式来构建一个后缀自动机然后在上面计数就好了。

那么如何计数呢？其实是一个传统的套路，在我们建完所有节点之后，直接统计每个节点的 \(maxlen_x - maxlen_{fa}\) . 我们考虑这样为什么是对的，因为每个点代表了一个等价类，那么其内部最短的串长其实就是其父亲的串长+1，所以：

\(ans =\sum maxlen_x-minlen_x+1\)
\(=maxlen_x-(maxlen_{fa}+1)+1\)
\(=maxlen_x - maxlen_{fa}\)

然后这题就做完了

Code:

#include<bits/stdc++.h>
#define ll long long
const int N=2e6+5;
using namespace std;
int n,m;
int col[N];
struct Trie{
    int ch[N][10],col[N],fa[N];
    int cnt;
    void init(){cnt=1;}
    int insert(int c,int p)
    {
        if(!ch[p][c])ch[p][c]=++cnt,fa[cnt]=p,col[cnt]=c;
        //cout<<p<<" "<<ch[p][c]<<"="<<c<<"\n";
        return ch[p][c];
    }
}T;
struct SAM{
    int ch[N<<1][10],len[N<<1],fa[N<<1];
    int cnt;void init(){cnt=1;}
    int insert(int c,int last)
    {
        int p=last,q=++cnt;len[q]=len[p]+1;
        for(;p&&!ch[p][c];p=fa[p])ch[p][c]=q;
        if(!p){fa[q]=1;return q;}int x=ch[p][c];
        if(len[x]==len[p]+1){fa[q]=x;return q;}
        int y=++cnt;len[y]=len[p]+1;fa[y]=fa[x];
        for(int i=0;i<m;i++)ch[y][i]=ch[x][i];
        for(;p&&ch[p][c]==x;p=fa[p])ch[p][c]=y;
        fa[x]=fa[q]=y;return q;
    }
    queue<int> Q;
    int pos[N<<1];
    void build()
    {
        for(int i=0;i<m;i++)if(T.ch[1][i])Q.push(T.ch[1][i]);
        pos[1]=1;
        while(!Q.empty())
        {
            int x=Q.front();Q.pop();
            pos[x]=insert(T.col[x],pos[T.fa[x]]);
            for(int i=0;i<m;i++)if(T.ch[x][i])Q.push(T.ch[x][i]);
        }
    }
    ll calc()
    {
        ll res=0;for(int u=2;u<=cnt;u++)res+=len[u]-len[fa[u]];
        return res;
    }
}sam;
struct Edge{
    int to,nxt;
}e[N];int head[N],d[N];
inline void add(int x,int y){e[++head[0]]={y,head[x]};head[x]=head[0];}
void dfs(int x,int fa,int fa_pos)
{
    int x_pos=T.insert(col[x],fa_pos);
    for(int i=head[x],y;i;i=e[i].nxt)
    if((y=e[i].to)!=fa){dfs(y,x,x_pos);}
}
void work()
{
    cin>>n>>m;
    sam.init();T.init();
    for(int i=1;i<=n;i++)scanf("%d",&col[i]);
    for(int i=1,x,y;i<n;i++)
    {
        scanf("%d%d",&x,&y);
        ++d[x];add(x,y);
        ++d[y];add(y,x);
    }
    for(int i=1;i<=n;i++)
    {
        if(d[i]==1)dfs(i,0,1);
    }
    sam.build();
    ll ans=sam.calc();
    printf("%lld",ans);
}
int main()
{
    //freopen("substring.in","r",stdin); freopen("substring.out","w",stdout);
    work();
    return 0;
}