C++14(g++5.4) 解法, 执行用时: 193ms, 内存消耗: 232256K, 提交时间: 2020-07-20 02:17:53

#include<bits/stdc++.h>
using namespace std;
const int N=2e6+100;
typedef long long ll;

struct SAM{
    int s[N];
    int ch[N*2][26],len[N*2],fa[N*2],a[N],n;
    int tot=1,last=1;
    int endpos_size[N],c[N];//
    vector<int>edge[N];

    int add(int c,int last) {
        if(ch[last][c]) {
            int p=last,x=ch[p][c];
            if(len[p]+1==len[x]) return x;
            else {
                int y=++tot;
                len[y]=len[p]+1;
                memcpy(ch[y],ch[x],sizeof(ch[x]));
                while(p&&ch[p][c]==x) ch[p][c]=y,p=fa[p];
                fa[y]=fa[x];fa[x]=y;
                return y;
            }
        }
        int p=last,np=++tot;
        len[np]=len[p]+1; endpos_size[np]=1;
        for(;p&&!ch[p][c];p=fa[p]) ch[p][c]=np;
        if(!p) fa[np]=1;
        else {
            int q=ch[p][c];
            if(len[q]==len[p]+1) fa[np]=q;
            else {
                int nq=++tot;len[nq]=len[p]+1;
                memcpy(ch[nq],ch[q],sizeof(ch[q]));
                fa[nq]=fa[q];fa[q]=fa[np]=nq;
                for(;p&&ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
            }
        }
        return np;
    }
    ll getans() {
        ll ans=0;
        for(int i=2;i<=tot;i++) ans+=len[i]-len[fa[i]];
        return ans;
    }
    void dfs(int x,int fa,int st) {
        st=add(s[x],st);
        for(auto v:edge[x]) if(v!=fa)
            dfs(v,x,st);
    }
    int deg[N];
    void solve() {
        cin>>n;int c;cin>>c;
        for(int i=1;i<=n;i++) scanf("%d",&s[i]);
        for(int i=1;i<n;i++) {
            int u,v;scanf("%d%d",&u,&v);
            edge[u].push_back(v);
            edge[v].push_back(u);
            deg[u]++;deg[v]++;
        }
        for(int i=1;i<=n;i++) if(deg[i]==1) dfs(i,0,1);
        cout<<getans();
    }
}sam;


int main() {
    sam.solve();
}

C++ 解法, 执行用时: 167ms, 内存消耗: 83064K, 提交时间: 2021-08-20 17:21:34

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<queue>
#include<cstring>
#include<vector>
using namespace std;

int nxt[4000010][10],gs,len[4000010],link[4000010];
int n,m,col[100010];
vector<int> v[100010]; 

int insert(int x,int c){
	if(nxt[x][c])return nxt[x][c];
	return nxt[x][c]=++gs;
}

int update(int lst,int c){
	int dq=nxt[lst][c];len[dq]=len[lst]+1;
	int p=link[lst];
	while(p!=-1 && !nxt[p][c]){
		nxt[p][c]=dq;p=link[p];
	}
	if(p==-1)return dq;
	int q=nxt[p][c];
	if(len[q]==len[p]+1){
		link[dq]=q;return dq;
	}
	int clo=++gs;
	for(int i=0;i<m;i++)nxt[clo][i]=len[nxt[q][i]]?nxt[q][i]:0;
	len[clo]=len[p]+1;
	while(p!=-1 && nxt[p][c]==q){
		nxt[p][c]=clo;p=link[p];
	}
	link[clo]=link[q];link[q]=clo;link[dq]=clo;return dq;
}

void dfs(int x,int lst,int fa){
	int dq=insert(lst,col[x]);
	for(int i:v[x]){
		if(i==fa)continue;
		dfs(i,dq,x);
	}
}

queue<pair<int,int> >q;

int main(){
	scanf("%d %d",&n,&m);link[0]=-1;
	for(int i=1;i<=n;i++)scanf("%d",&col[i]);
	for(int i=1;i<n;i++){
		int x,y;
		scanf("%d %d",&x,&y);
		v[x].push_back(y);v[y].push_back(x);
	}
	for(int i=1;i<=n;i++){
		if(v[i].size()==1)dfs(i,0,0);
	}
	for(int i=0;i<m;i++)if(nxt[0][i])q.push(make_pair(0,i));
	while(!q.empty()){
		auto dq=q.front();q.pop();
		int lst=update(dq.first,dq.second);
		for(int i=0;i<m;i++)if(nxt[lst][i])q.push(make_pair(lst,i));
	}
	long long ans=0;
	for(int i=1;i<=gs;i++)ans+=len[i]-len[link[i]];
	printf("%lld\n",ans);
	return 0;
}