Codeforces 620E New Year Tree（DFS序+线段树）

  1 #include<cstdio>
  2 #include<cstring>
  3 using namespace std;
  4 #define MAXN 500000
  5 struct Edge{
  6     int v,next;
  7 }edge[MAXN<<1];
  8 int NE,head[MAXN];
  9 void addEdge(int u,int v){
 10     edge[NE].v=v; edge[NE].next=head[u];
 11     head[u]=NE++;
 12 }
 13 
 14 int l[MAXN],r[MAXN],dfn;
 15 void dfs(int u,int fa){
 16     l[u]=++dfn;
 17     for(int i=head[u]; i!=-1; i=edge[i].next){
 18         int v=edge[i].v;
 19         if(v==fa) continue;
 20         dfs(v,u);
 21     }
 22     r[u]=dfn;
 23 }
 24 
 25 long long tree[MAXN<<2],lazy[MAXN<<2],z;
 26 int N,x,y;
 27 void update(int i,int j,int k){
 28     if(x<=i && j<=y){
 29         tree[k]=z;
 30         lazy[k]=z;
 31         return;
 32     }
 33     if(lazy[k]){
 34         tree[k<<1]=lazy[k];
 35         tree[k<<1|1]=lazy[k];
 36         lazy[k<<1]=lazy[k];
 37         lazy[k<<1|1]=lazy[k];
 38         lazy[k]=0;
 39     }
 40     int mid=i+j>>1;
 41     if(x<=mid) update(i,mid,k<<1);
 42     if(y>mid) update(mid+1,j,k<<1|1);
 43     tree[k]=tree[k<<1]|tree[k<<1|1];
 44 }
 45 long long query(int i,int j,int k){
 46     if(x<=i && j<=y){
 47         return tree[k];
 48     }
 49     if(lazy[k]){
 50         tree[k<<1]=lazy[k];
 51         tree[k<<1|1]=lazy[k];
 52         lazy[k<<1]=lazy[k];
 53         lazy[k<<1|1]=lazy[k];
 54         lazy[k]=0;
 55     }
 56     int mid=i+j>>1;
 57     long long res=0;
 58     if(x<=mid) res|=query(i,mid,k<<1);
 59     if(y>mid) res|=query(mid+1,j,k<<1|1);
 60     return res;
 61 }
 62 
 63 int getCount(long long s){
 64     int cnt=0;
 65     for(int i=0; i<60; ++i){
 66         if(s>>i&1) ++cnt;
 67     }
 68     return cnt;
 69 }
 70 
 71 int color[MAXN];
 72 int main(){
 73     int n,m;
 74     scanf("%d%d",&n,&m);
 75     for(int i=1; i<=n; ++i){
 76         scanf("%d",color+i);
 77         --color[i];
 78     }
 79     memset(head,-1,sizeof(head));
 80     int a,b,c;
 81     for(int i=1; i<n; ++i){
 82         scanf("%d%d",&a,&b);
 83         addEdge(a,b);
 84         addEdge(b,a);
 85     }
 86     dfs(1,1);
 87     for(N=1; N<n; N<<=1);
 88     for(int i=1; i<=n; ++i){
 89         x=l[i]; y=l[i]; z=1LL<<color[i];
 90         update(1,N,1);
 91     }
 92     while(m--){
 93         scanf("%d",&c);
 94         if(c==1){
 95             scanf("%d%d",&a,&b);
 96             x=l[a]; y=r[a]; z=1LL<<b-1;
 97             update(1,N,1);
 98         }else{
 99             scanf("%d",&a);
100             x=l[a]; y=r[a];
101             printf("%d\n",getCount(query(1,N,1)));
102         }
103     }
104     return 0;
105 }