BZOJ2243 [SDOI2011] 染色

  1 #include <iostream>
  2 #include <cstdio>
  3 #include <algorithm>
  4 #include <cstring>
  5 #define rep(i,l,r) for(int i=l; i<=r; i++)
  6 #define clr(x,y) memset(x,y,sizeof(x))
  7 #define travel(x) for(Edge *p=last[x]; p; p=p->pre)
  8 using namespace std;
  9 const int maxn = 100010;
 10 struct Edge{
 11     Edge *pre;
 12     int to;
 13 }edge[maxn<<1];
 14 Edge *last[maxn],*pt;
 15 struct node{
 16     int l,r,c,s,lc,rc;
 17 }t[maxn<<2];
 18 int n,m,x,y,z,segnum=0,c[maxn],depth[maxn],fa[maxn],pos[maxn],size[maxn],belong[maxn];
 19 char op[10];
 20 bool vis[maxn];
 21 inline int read(){
 22     int ans = 0, f = 1;
 23     char c = getchar();
 24     while (!isdigit(c)){
 25         if (c == ‘-‘) f = -1;
 26         c = getchar();
 27     }
 28     while (isdigit(c)){
 29         ans = ans * 10 + c - ‘0‘;
 30         c = getchar();
 31     }
 32     return ans * f;
 33 }
 34 inline void addedge(int x,int y){
 35     pt->pre = last[x];
 36     pt->to = y;
 37     last[x] = pt++;
 38 }
 39 void dfs1(int x){
 40     vis[x] = 1; size[x] = 1;
 41     travel(x){
 42         if (vis[p->to]) continue;
 43         depth[p->to] = depth[x] + 1;
 44         fa[p->to] = x;
 45         dfs1(p->to);
 46         size[x] += size[p->to];
 47     }
 48 }
 49 void dfs2(int x,int chain){
 50     int k = 0; pos[x] = ++segnum; belong[x] = chain;
 51     travel(x){
 52         if (depth[p->to] > depth[x] && size[p->to] > size[k])
 53         k = p->to;
 54     }
 55     if (!k) return;
 56     dfs2(k,chain);
 57     travel(x){
 58         if (depth[p->to] > depth[x] && p->to != k) dfs2(p->to,p->to);
 59     }
 60 }
 61 inline void pushdown(int w){
 62     if (t[w].c == -1 || t[w].l == t[w].r) return;
 63     t[w<<1].s = t[w<<1|1].s = 1;
 64     t[w<<1].c = t[w<<1|1].c = t[w].c;
 65     t[w<<1].lc = t[w<<1|1].lc = t[w<<1].rc = t[w<<1|1].rc = t[w].c;
 66     t[w].c = -1;
 67 }
 68 inline void maintain(int w){
 69     t[w].lc = t[w<<1].lc; t[w].rc = t[w<<1|1].rc;
 70     t[w].s = t[w<<1].s + t[w<<1|1].s - (t[w<<1].rc == t[w<<1|1].lc);
 71 }
 72 void build(int u,int v,int w){
 73     t[w].l = u; t[w].r = v; t[w].c = -1; t[w].s = 1;
 74     if (u == v) return;
 75     int mid = (u + v) >> 1;
 76     build(u,mid,w<<1); build(mid+1,v,w<<1|1);
 77 }
 78 void change(int u,int v,int w,int x){
 79     pushdown(w);
 80     if (u == t[w].l && v == t[w].r){
 81         t[w].c = t[w].lc = t[w].rc = x; 
 82         t[w].s = 1; return;
 83     }
 84     int mid = (t[w].l + t[w].r) >> 1;
 85     if (v <= mid) change(u,v,w<<1,x);
 86     else if (u > mid) change(u,v,w<<1|1,x);
 87     else{
 88         change(u,mid,w<<1,x);
 89         change(mid+1,v,w<<1|1,x);
 90     }
 91     maintain(w);
 92 }
 93 int queryseg(int u,int v,int w){
 94     pushdown(w);
 95     if (u == t[w].l && v == t[w].r) return t[w].s;
 96     int mid = (t[w].l + t[w].r) >> 1;
 97     if (v <= mid) return queryseg(u,v,w<<1);
 98     else if (u > mid) return queryseg(u,v,w<<1|1);
 99     else{
100         int ret = queryseg(u,mid,w<<1) + queryseg(mid+1,v,w<<1|1);
101         if (t[w<<1].rc == t[w<<1|1].lc) ret--;
102         return ret;
103     }
104 }
105 int find(int x,int w){
106     pushdown(w);
107     if (t[w].l == t[w].r) return t[w].lc;
108     int mid = (t[w].l + t[w].r) >> 1;
109     if (x <= mid) return find(x,w<<1);
110     else return find(x,w<<1|1);
111 }
112 void modify(int x,int y,int z){
113     while (belong[x] != belong[y]){
114         if (depth[belong[x]] < depth[belong[y]]) swap(x,y);
115         change(pos[belong[x]],pos[x],1,z);
116         x = fa[belong[x]];
117     }
118     if (depth[x] < depth[y]) swap(x,y);
119     change(pos[y],pos[x],1,z);
120 }
121 int query(int x,int y){
122     int ret = 0;
123     while (belong[x] != belong[y]){
124         if (depth[belong[x]] < depth[belong[y]]) swap(x,y);
125         ret += queryseg(pos[belong[x]],pos[x],1);
126         if (find(pos[belong[x]],1) == find(pos[fa[belong[x]]],1)) ret--;
127         x = fa[belong[x]];
128     }
129     if (depth[x] < depth[y]) swap(x,y);
130     ret += queryseg(pos[y],pos[x],1);
131     return ret;
132 }
133 int main(){
134     n = read(); m = read(); clr(last,0); pt = edge;
135     rep(i,1,n) c[i] = read();
136     rep(i,1,n-1){
137         x = read(); y = read();
138         addedge(x,y); addedge(y,x);
139     }
140     depth[1] = 0; dfs1(1); dfs2(1,1); build(1,n,1);
141     rep(i,1,n) change(pos[i],pos[i],1,c[i]);
142     rep(i,1,m){
143         scanf("%s",op);
144         switch(op[0]){
145             case ‘C‘:
146                 x = read(); y = read(); z = read();
147                 modify(x,y,z);
148             break;
149             case ‘Q‘:
150                 x = read(); y = read();
151                 printf("%d\n",query(x,y));
152             break;
153         }
154     }
155     return 0;
156 }