线段树动态开点+树链剖分BZOJ4999

以每个一个颜色开一颗线段树，内部以dfs序作为线段树节点，权值代表出现次数，维护线段树区间和

#include<iostream>
#include<stdio.h>
#include<string.h>
#include<algorithm>
#include<map>
using namespace std;
const int maxn = 2e5+6;
struct edge{
   int next,to;
}e[maxn*2];
struct node{
  int l,r;
  int w;
}tree[maxn*50];
int head[maxn],siz[maxn],top[maxn],son[maxn];
int d[maxn],fa[maxn],id[maxn],rk[maxn],dfn[maxn],c[maxn];
int n,m,uu,vv,tot,cnt;
int rt[maxn];
map<int,int>p;
void add(int x,int y){
    e[++cnt].next=head[x];
    e[cnt].to=y;
    head[x]=cnt;
}
void dfs1(int u,int f,int depth){
   d[u]=depth;
   fa[u]=f;
   siz[u]=1;
   for (int i=head[u];i;i=e[i].next){
      int v=e[i].to;
      if (v==f)continue;
      dfs1(v,u,depth+1);
      siz[u]+=siz[v];
      if (siz[v]>siz[son[u]])
        son[u]=v;
   }
}
void dfs2(int u,int t){
   top[u]=t;
   dfn[u]=++cnt;
   if (!son[u])
     return;
   dfs2(son[u],t);
   for (int i=head[u];i;i=e[i].next)
   {
       int v=e[i].to;
       if (v!=son[u] && v!=fa[u])
         dfs2(v,v);
   }
}
void Modify(int &x,int l,int r,int pos,int w)
{
    if (!x)x=++tot; tree[x].w+=w;
    if (l==r)return;
    int mid=(l+r)>>1;
    if (pos<=mid)Modify(tree[x].l,l,mid,pos,w);
    else Modify(tree[x].r,mid+1,r,pos,w);
}
int query(int x,int l,int r,int L,int R){
    if((!tree[x].w) || (!x))return 0;
    if  (L<=l && r<=R){
         return tree[x].w;
    }
    int mid=(l+r)>>1;
    int ans=0;
    if (L<=mid)ans+=query(tree[x].l,l,mid,L,R);
    if (R>mid)ans+=query(tree[x].r,mid+1,r,L,R);
    return ans;
}
int qryRange(int x,int y,int k){
   int res=0;
   while(top[x]!=top[y]){
      if (d[top[x]]<d[top[y]])swap(x,y);
      res+=query(rt[k],1,n,dfn[top[x]],dfn[x]);
      x=fa[top[x]];
   }
   if (d[y]>d[x])swap(x,y);
   res+=query(rt[k],1,n,dfn[y],dfn[x]);
   return res;
}
int main(){
  int sum=0,q;
  while(~scanf("%d%d",&n,&q)){
      cnt=0;
      sum=0;
      for (int i=1;i<=n;i++){
         scanf("%d",&c[i]);
      }
      for (int i=1;i<=n;i++){
        scanf("%d%d",&uu,&vv);
        add(uu,vv);
        add(vv,uu);
      }
      cnt=0;
      dfs1(1,0,1);
      dfs2(1,1);
      tot=0;
      for (int i=1;i<=n;i++){
        if (p[c[i]]==0){
            p[c[i]]=++sum;
        }
        Modify(rt[p[c[i]]],1,n,dfn[i],1);
      }
      char tmp[3];
      int u,w,v;
      while(q--){
         scanf("%s",tmp);
         if (tmp[0]=='C')
         {
             scanf("%d%d",&u,&w);
             Modify(rt[p[c[u]]],1,n,dfn[u],-1);
             if (!p[w])p[w]=++sum;
             Modify(rt[p[w]],1,n,dfn[u],1);
             c[u]=w;
         }else {
             scanf("%d%d%d",&u,&v,&w);
             if (!p[w])printf("0\n");
             else printf("%d\n",qryRange(u,v,p[w]));
         }
      }
  }
  return 0;
}