以每个一个颜色开一颗线段树,内部以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 "); else printf("%d ",qryRange(u,v,p[w])); } } } return 0; }