BZOJ3926: [Zjoi20150]诸神眷顾的幻想乡

这一题有点坑。。注意longlong

SAM裸题。。。

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



struct Node
{
    int len;
    Node *f,*last,*ch[10];
    Node(){for(len=0;len<=9;len++)ch[len]=NULL;last=NULL;len=0;}
}*root,*last;

inline void add(int data)
{
    Node *tp,*ne=new Node;
    ne->len=last->len+1;
    ne->f=last;
    for(tp=last;tp&&!tp->ch[data];tp=tp->last)tp->ch[data]=ne;
    if(!tp)
      ne->last=root,last=ne;
   else if(tp->len==tp->ch[data]->len-1)
       ne->last=tp->ch[data],last=ne;
    else
      {
        Node *a=new Node,*b=tp->ch[data];
        *a=*b;
        ne->last=b->last=a;
        a->len=tp->len+1;
        last=ne;
        a->f=tp;
        for(tp;tp&&tp->ch[data]==b;tp=tp->last)tp->ch[data]=a;
      }
}
struct Chain
{
    Chain*next;
    int u;
    Chain(){next=NULL;}
}*Head[1000001];

inline void addside(int a,int b)
{
  Chain *tp=new Chain;
  tp->u=b,tp->next=Head[a],Head[a]=tp;
}
int color[100001];
void DFS(int u,int fa)
{
    for(Chain *tp=Head[u];tp;tp=tp->next)
      if(fa^tp->u)
          add(color[tp->u]),DFS(tp->u,u);
    last=last->f;
}

int n,C;
char c;
inline void read(int &a)
{
    a=0;do c=getchar();while(c<'0'||c>'9');
    while(c<='9'&&c>='0')a=(a<<3)+(a<<1)+c-'0',c=getchar();
}
int Sort[1000001];
int H[1000001];
inline bool cmp(int x,int y)
{return H[x]<H[y];} 
long long ans;
void Solve(Node *tp)
{
    if(tp->last)
       ans+=tp->len-tp->last->len;
    for(int i=0;i<=9;i++)
      if(tp->ch[i]&&tp->ch[i]->len==tp->len+1)
         Solve(tp->ch[i]);
}
int main()
{
    read(n),read(C);
     int i,j,k;
     for(i=1;i<=n;i++)
     Sort[i]=i,read(color[i]);
    root=last=new Node;
     for(k=1;k<n;k++)
        read(i),read(j),addside(i,j),addside(j,i),H[i]++,H[j]++;
    sort(Sort+1,Sort+1+n,cmp);
    for(i=1;i<=n;i++)
       if(H[Sort[i]]==1)
          {
            last=root;
            add(color[Sort[i]]);
            DFS(Sort[i],0);
          }
        else break;
    Solve(root);
    printf("%lld\n",ans);
    return 0;
}