记录编号 272719 评测结果 AAAAAAAAAAAAAAAAAAAA
题目名称 [ZJOI 2015]诸神眷顾的幻想乡 最终得分 100
用户昵称 GravatarTenderRun 是否通过 通过
代码语言 C++ 运行时间 2.496 s
提交时间 2016-06-17 20:06:43 内存使用 182.59 MiB
显示代码纯文本
#include <iostream>
#include <cstring>
#include <cstdio>
using namespace std;
const int maxn=100010;
const int maxm=2000010;
long long ans;
bool vis[maxm];
int n,sumc,cntE,c[maxn];
int fir[maxn],nxt[maxn<<1],to[maxn<<1];
void addedge(int a,int b){
	nxt[++cntE]=fir[a];
	fir[a]=cntE;
	to[cntE]=b;
}
int st[maxm],topst;
int CH[maxm][11],pos[maxm];
int fa[maxm],ch[maxm][11],len[maxm];
struct SAM{
	int cnt,last;
	SAM(){
		cnt=last=1;
	}
	int Insert(int c){
		int p=last,np=last=++cnt;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[p]+1==len[q])
				fa[np]=q;
			else{
				int nq=++cnt;len[nq]=len[p]+1;
				memcpy(ch[nq],ch[q],sizeof(ch[q]));
				fa[nq]=fa[q];fa[q]=fa[np]=nq;
				while(ch[p][c]==q)ch[p][c]=nq,p=fa[p];
			}	
		}
		ans+=len[last]-len[fa[last]];
		return last;
	}
}sam;
struct Trie{
	int cnt,rt;
	Trie(){
		rt=cnt=1;
		pos[rt]=1;
	}
	void Insert(int &x,int p){
		vis[p]=true;
		if(!x)x=++cnt;
		for(int i=fir[p];i;i=nxt[i])
			if(!vis[to[i]])
				Insert(CH[x][c[to[i]]],to[i]);
	}
	void Construct(int x){
		for(int i=0;i<sumc;i++)
			if(CH[x][i]){
				sam.last=pos[x];
				pos[CH[x][i]]=sam.Insert(i);
			}
		for(int i=0;i<sumc;i++)
			if(CH[x][i])
				Construct(CH[x][i]);
	}
}trie;
int main(){
#ifndef ONLINE_JUDGE
	freopen("zjoi15_substring.in","r",stdin);
	freopen("zjoi15_substring.out","w",stdout);
#endif
	scanf("%d%d",&n,&sumc);
	for(int i=1;i<=n;i++)
		scanf("%d",&c[i]); 
	for(int i=1,a,b;i<n;i++){
		scanf("%d%d",&a,&b);
		addedge(a,b);
		addedge(b,a);
	}	
	for(int x=1,tot=0;x<=n;tot=0,x++){
		for(int i=fir[x];i;i=nxt[i])tot+=1;
		if(tot==1)st[++topst]=x;
	}
	while(topst){
		memset(vis,0,sizeof(vis));
		trie.Insert(CH[trie.rt][c[st[topst]]],st[topst]);
		topst-=1;
	}
	trie.Construct(trie.rt);
	printf("%lld\n",ans);
	return 0;	
}