[BZOJ2243][SDOI2011]染色

  1 #include<iostream>
  2 #include<cstdio>
  3 #include<cstring>
  4 #include<cstdlib>
  5 #include<cmath>
  6 #include<algorithm>
  7 using namespace std;
  8 int n,m;
  9 int c[400005];
 10 struct data
 11 {
 12     int next,to;
 13 }e[400005];
 14 int head[400005],cnt;
 15 void add(int u,int v){e[cnt].next=head[u];e[cnt].to=v;head[u]=cnt;cnt++;}
 16 int son[400005],size[400005],fa[400005],dep[400005];
 17 void findson(int now)
 18 {
 19     size[now]=1;
 20     for(int i=head[now];i>=0;i=e[i].next)
 21     {
 22         int to=e[i].to;
 23         if(to==fa[now]) continue;
 24         fa[to]=now;
 25         dep[to]=dep[now]+1;
 26         findson(to);
 27         size[now]+=size[to];
 28         if(size[to]>size[son[now]]) son[now]=to;
 29     }
 30 }
 31 int up[400005],id[400005],p[400005],s;
 32 void dfs(int now,int top)
 33 {
 34     id[now]=++s;
 35     p[s]=now;
 36     up[now]=top;
 37     if(son[now]) dfs(son[now],top);
 38     for(int i=head[now];i>=0;i=e[i].next)
 39     {
 40         int to=e[i].to;
 41         if(to==fa[now]||to==son[now]) continue;
 42         dfs(to,to);
 43     }
 44 }
 45 struct te
 46 {
 47     int lc,rc,color,l,r,lazy;
 48 }t[800005];
 49 void pushup(int now)
 50 {
 51     int l=now<<1,r=l|1;
 52     t[now].lc=t[l].lc;t[now].rc=t[r].rc;
 53     t[now].color=t[l].color+t[r].color;
 54     if(t[l].rc==t[r].lc) t[now].color--;
 55 }
 56 void pushdown(int now)
 57 {
 58     if(!t[now].lazy) return;
 59     int l=now<<1,r=l|1;
 60     t[l].lazy=t[l].lc=t[l].rc=t[r].lazy=t[r].lc=t[r].rc=t[now].lazy;
 61     t[l].color=t[r].color=1;
 62     t[now].lazy=0;
 63 }
 64 void build(int now,int l,int r)
 65 {
 66     t[now].l=l;t[now].r=r;
 67     if(l==r){t[now].color=1;t[now].lc=t[now].rc=c[p[l]];return;}
 68     else
 69     {
 70         int mid=(l+r)>>1;
 71         build(now<<1,l,mid);
 72         build(now<<1|1,mid+1,r);
 73     }
 74     pushup(now);
 75 }
 76 void update(int now,int L,int R,int c)
 77 {
 78     int l=t[now].l,r=t[now].r;
 79     pushdown(now);
 80     if(L<=l&&R>=r){t[now].lc=t[now].rc=t[now].lazy=c;t[now].color=1;return;}
 81     else
 82     {
 83         int mid=(l+r)>>1;
 84         if(L<=mid) update(now<<1,L,R,c);
 85         if(R>mid) update(now<<1|1,L,R,c);
 86     }
 87     pushup(now);
 88 }
 89 int findlca(int a,int b)
 90 {
 91     int x,y;
 92     x=up[a],y=up[b];
 93     while(x!=y)
 94     {
 95         if(dep[x]<dep[y]) swap(x,y),swap(a,b);
 96         a=fa[x];x=up[a];
 97     }
 98     if(dep[a]<dep[b]) swap(a,b);
 99     return b;
100 }
101 int query(int now,int L,int R)
102 {
103     int ans=0;
104     int l=t[now].l,r=t[now].r;
105     pushdown(now);
106     if(L<=l&&R>=r) return t[now].color;
107     else
108     {
109            int mid=(l+r)>>1;
110            if(R<=mid) ans+=query(now<<1,L,R);
111            else if(L>mid) ans+=query(now<<1|1,L,R);
112            else
113            {
114             ans+=query(now<<1,L,R)+query(now<<1|1,L,R);
115             if(t[now<<1].rc==t[now<<1|1].lc) ans--;
116         }
117        }
118        return ans;
119 }
120 void change(int now,int lca,int c)
121 {
122     while(up[now]!=up[lca])
123     {
124            update(1,id[up[now]],id[now],c);
125            now=fa[up[now]];
126        }
127        update(1,id[lca],id[now],c);
128 }
129 int getc(int now,int x)
130 {
131     pushdown(now);
132     int l=t[now].l,r=t[now].r;
133     if(l==r) return t[now].lc;
134     int mid=(l+r)>>1;
135     if(x<=mid)return getc(now<<1,x);
136     else return getc(now<<1|1,x);
137 }
138 
139 int ask(int now,int lca)
140 {
141     int ans=0;
142     while(up[now]!=up[lca])
143     {
144            ans+=query(1,id[up[now]],id[now]);
145            if(getc(1,id[up[now]])==getc(1,id[fa[up[now]]])) ans--;
146            now=fa[up[now]];
147        }
148        return ans+query(1,id[lca],id[now]);
149 }
150 int main()
151 {
152     memset(head,-1,sizeof(head));
153     scanf("%d%d",&n,&m);
154     for(int i=1;i<=n;i++) scanf("%d",&c[i]);
155     for(int i=1;i<=n-1;i++)
156     {
157         int u,v;
158         scanf("%d%d",&u,&v);
159         add(u,v);add(v,u);
160     }
161     dep[1]=1;
162     findson(1);dfs(1,1);build(1,1,n);
163     for(int i=1;i<=m;i++)
164     {
165         char ch[10];
166         scanf("%s",ch);
167         if(ch[0]==‘C‘)
168         {
169             int a,b,c;
170             scanf("%d%d%d",&a,&b,&c);
171             int lca=findlca(a,b);
172             change(a,lca,c);change(b,lca,c);
173         }
174         else
175         {
176                int a,b;
177                scanf("%d%d",&a,&b);
178                int lca=findlca(a,b);
179                printf("%d\n",ask(a,lca)+ask(b,lca)-1);
180            }
181     }
182 }