[ZJOI2015] 诸神眷顾的幻想乡

P3256【ZJOI2015 Day1】诸神眷顾的幻想乡

时间限制 : 20000 MS   空间限制 : 524288 KB

问题描述

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

输入格式

        第一行两个正整数n,c。表示空地数量和颜色数量。
        第二行有n个0到c-1之间，由空格隔开的整数，依次表示第i块空地上的粉丝的衣服颜色。（这里我们按照节点标号从小到大的顺序依次给出每块空地上粉丝的衣服颜色）。
        接下来n-1行，每行两个正整数u,v，表示有一条连接空地u和空地v的边。

输出格式

        一行，输出一个整数，表示答案。

样例输入

7 3
0 2 1 2 1 0 0 
1 2
3 4
3 5
4 6
5 7
2 5

样例输出

30

提示

        对于所有数据，1<=n<=100000, 1<=c<=10。
        对于15%的数据，n<=2000。
        另有5%的数据，所有空地都至多与两个空地相邻。
        另有5%的数据，除一块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。
        另有5%的数据，除某两块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。

题解：

发现每一个子串均是以某一个叶子结点为根的树中，根到叶子结点的路径上的后缀的前缀，或者说是这条路径上的子串。

因为保证了叶子结点<=20个，那么，我们以每个叶子结点为根开始算，子串个数我们可以用后缀自动机来计算。

#include<cstdio> 
#include<iostream> 
#include<cstdlib> 
#include<algorithm> 
#include<cstring> 
#define ll long long 
using namespace std; 
template <typename T> 
inline void _read(T& x){ 
    char t=getchar();bool sign=true; 
    while(t<'0'||t>'9'){if(t=='-')sign=false;t=getchar();} 
    for(x=0;t>='0'&&t<='9';t=getchar())x=x*10+t-'0'; 
    if(!sign)x=-x; 
} 
int n,m,c,e,tot=1,last=1,root=1; 
struct node{
    int maxn; 
    int son[12]; 
};  
node _auto[2000005]; 
int par[2000005]; 
int insert(int t,int auto_fa){
    int np,p=auto_fa;
    np=last=++tot;
    _auto[np].maxn=_auto[p].maxn+1;
    while(!_auto[p].son[t]&&p)_auto[p].son[t]=np,p=par[p];
    if(!p)par[np]=1;
    else{
        int q=_auto[p].son[t];
        if(_auto[p].maxn+1==_auto[q].maxn)par[np]=q;
        else{
        	int nq=++tot;_auto[nq].maxn=_auto[p].maxn+1;
        	memcpy(_auto[nq].son,_auto[q].son,sizeof(_auto[q].son));
        	par[nq]=par[q];
        	par[np]=par[q]=nq;
        	while(_auto[p].son[t]==q)_auto[p].son[t]=nq,p=par[p];
        }
    }
    return np;
}

struct line{
	int from,to;
	line(){}
	line(int x,int y){from=x;to=y;}
};
line edge[200005];
int Last[100005],_next[200005];
void add_edge(int x,int y){
	edge[++e]=line(x,y);
	_next[e]=Last[x];
	Last[x]=e;
}
int color[100005],deg[100005];
void dfs(int x,int fa,int auto_fa){
	int temp=insert(color[x],auto_fa);
	for(int i=Last[x];i;i=_next[i]){
		int v=edge[i].to;
		if(v==fa)continue;
		dfs(v,x,temp);
	}
}
int main_main(){
	int i,j,k;
	cin>>n>>c;
	for(i=1;i<=n;i++){
		_read(color[i]);
	}
	for(i=1;i<n;i++){
		int x,y;
		_read(x);_read(y);
		add_edge(x,y);
		add_edge(y,x);
		deg[x]++;deg[y]++;
	}
	for(i=1;i<=n;i++){
		if(deg[i]==1)dfs(i,0,1);
	}
	ll ans=0;
	for(i=1;i<=tot;i++){
		ans+=_auto[i].maxn-_auto[par[i]].maxn;
	}
	cout<<ans;
} 
const int main_stack=16;  
char my_stack[128<<20];  
int main() {  
  __asm__("movl %%esp, (%%eax);\n"::"a"(my_stack):"memory");  
  __asm__("movl %%eax, %%esp;\n"::"a"(my_stack+sizeof(my_stack)-main_stack):"%esp");  
  main_main();  
  __asm__("movl (%%eax), %%esp;\n"::"a"(my_stack):"%esp");  
  return 0;  
}