【树链剖分】bzoj2243 [SDOI2011]染色

  1 #include<cstdio>
  2 #include<algorithm>
  3 using namespace std;
  4 #define lson rt<<1,l,m
  5 #define rson rt<<1|1,m+1,r
  6 #define maxn 100001
  7 int col[maxn<<2],L[maxn<<2],R[maxn<<2],delta[maxn<<2],n,m,Val[maxn],Num[maxn];
  8 int V[maxn<<1],First[maxn<<1],Next[maxn<<1],All_Lcol,All_Rcol,en;
  9 bool vis[maxn];
 10 int top[maxn],son[maxn],tot,dep[maxn],fa[maxn],siz[maxn],Map[maxn];
 11 inline void AddEdge(const int &UU,const int &VV){V[++en]=VV;Next[en]=First[UU];First[UU]=en;}
 12 inline void pushup(const int &rt)
 13 {
 14     L[rt]=L[rt<<1];
 15     R[rt]=R[rt<<1|1];
 16     col[rt]=col[rt<<1]+col[rt<<1|1]-(R[rt<<1]==L[rt<<1|1]);
 17 }
 18 inline void pushdown(const int &rt)
 19 {
 20     if(delta[rt]!=-1)
 21       {
 22           L[rt<<1]=R[rt<<1]=L[rt<<1|1]=R[rt<<1|1]=delta[rt];
 23           delta[rt<<1]=delta[rt<<1|1]=delta[rt];
 24           delta[rt]=-1;
 25           col[rt<<1]=col[rt<<1|1]=1;
 26       }
 27 }
 28 void buildtree(int rt,int l,int r)
 29 {
 30     delta[rt]=-1;
 31     if(l==r)
 32       {
 33           col[rt]=1;
 34           L[rt]=R[rt]=Val[Map[l]];//l是线段树中的新位置，
 35                                 //Map[l]表示新位置对应的在树中的编号是谁 
 36         return;
 37       }
 38     int m=l+r>>1;
 39     buildtree(lson);
 40     buildtree(rson);
 41     pushup(rt);
 42 }
 43 void update(int ql,int qr,int v,int rt,int l,int r)
 44 {
 45     if(ql<=l&&r<=qr)
 46       {
 47         delta[rt]=L[rt]=R[rt]=v;//更新当前结点的标记值 
 48         col[rt]=1;
 49         return;
 50       }
 51     pushdown(rt);//将该节点的标记下传到孩子们 
 52     int m=l+r>>1;
 53     if(ql<=m)
 54       update(ql,qr,v,lson);
 55     if(m<qr)
 56       update(ql,qr,v,rson);
 57     pushup(rt);
 58 }
 59 int query(int ql,int qr,int rt,int l,int r)
 60 {
 61     if(ql==l)
 62       All_Lcol=L[rt];
 63     if(r==qr)
 64       All_Rcol=R[rt];
 65     if(ql<=l&&r<=qr)
 66       return col[rt];
 67     pushdown(rt);//将该节点的标记下传到孩子们 
 68     int m=l+r>>1,res=0,now1=-1,now2=-1;
 69     if(ql<=m)
 70       {
 71           res+=query(ql,qr,lson);
 72           now1=R[rt<<1];
 73       }
 74     if(m<qr)
 75       {
 76           res+=query(ql,qr,rson);
 77           now2=L[rt<<1|1];
 78       }
 79     res-=(now1==now2 && now1!=-1);
 80     return res;
 81 }
 82 inline void Update(int u,int v,int val)
 83 {
 84     int f1=top[u],f2=top[v];
 85     while(f1!=f2)
 86       {
 87           if(dep[f1]<dep[f2])
 88             {
 89                 swap(u,v);
 90                 swap(f1,f2);
 91             }
 92           update(Num[f1],Num[u],val,1,1,n);
 93           u=fa[f1];
 94           f1=top[u];
 95       }
 96     if(dep[u]<dep[v])
 97       swap(u,v);
 98     update(Num[v],Num[u],val,1,1,n);
 99 }
100 inline int Query(int u,int v)
101 {
102     int f1=top[u],f2=top[v],ans=0,t1=-1,t2=-1;
103     while(f1!=f2)
104       {
105           if(dep[f1]<dep[f2])
106             {
107                 swap(u,v);
108                 swap(f1,f2);
109                 swap(t1,t2);
110             }
111           ans+=query(Num[f1],Num[u],1,1,n);
112           ans-=(t1==All_Rcol);
113           t1=All_Lcol;
114           u=fa[f1];
115           f1=top[u];
116       }
117     if(dep[u]<dep[v])
118       {
119           swap(u,v);
120           swap(t1,t2);
121       }
122     ans+=query(Num[v],Num[u],1,1,n);
123     ans-=((t1==All_Rcol&&t1!=-1)+(t2==All_Lcol&&t2!=-1));
124     return ans;
125 }
126 void dfs1(int cur,int father,int depth)
127 {
128     fa[cur]=father;
129     dep[cur]=depth;
130     siz[cur]=1;
131     for(int i=First[cur];i;i=Next[i])
132       if(!vis[V[i]])
133         {
134           vis[V[i]]=true;
135           dfs1(V[i],cur,depth+1);
136           siz[cur]+=siz[V[i]];
137           if(siz[V[i]]>siz[son[cur]])
138             son[cur]=V[i];
139           vis[V[i]]=false;
140         }
141 }
142 void dfs2(int cur)
143 {
144     if(son[cur]&&!vis[son[cur]])
145       {
146           vis[son[cur]]=true;
147           top[son[cur]]=top[cur];
148           Num[son[cur]]=++tot;
149           Map[tot]=son[cur];
150           dfs2(son[cur]);
151           vis[son[cur]]=false;
152       }
153     for(int i=First[cur];i;i=Next[i])
154       if(son[cur]!=V[i]&&!vis[V[i]])
155         {
156           vis[V[i]]=true;
157           top[V[i]]=V[i];
158           Num[V[i]]=++tot;
159           Map[tot]=V[i];
160           dfs2(V[i]);
161           vis[V[i]]=false;
162         }
163 }
164 int Res,num;
165 char C;
166 char CH[20];
167 inline int getint()
168 {
169     Res=0;
170     C=‘*‘; 
171     while(C<‘0‘||C>‘9‘)
172       C=getchar();
173     while(C>=‘0‘&&C<=‘9‘)
174       {
175         Res=Res*10+(C-‘0‘);
176         C=getchar();
177       }
178     return Res;
179 }
180 inline void putint(int x)
181 {
182     num=0;  
183     while(x>0)
184       CH[++num]=x%10,x/=10;
185     while(num)   
186       putchar(CH[num--]+48);
187     putchar(‘\n‘);
188 }
189 int main()
190 {
191     int x,y,a,b,c;
192     char op;
193     n=getint();
194     m=getint();
195     for(int i=1;i<=n;i++)
196       Val[i]=getint();
197     for(int i=1;i<n;i++){x=getint();y=getint();AddEdge(x,y);AddEdge(y,x);}
198     top[1]=1;
199     Num[1]=++tot;
200     Map[tot]=1;
201     vis[1]=true;
202     dfs1(1,0,1);
203     dfs2(1);
204     buildtree(1,1,n);
205     for(int i=1;i<=m;i++)
206       {
207           getchar();
208           op=getchar();
209           a=getint();
210           b=getint();
211           if(op==‘C‘)
212             {
213                 c=getint();
214                 Update(a,b,c);
215             }
216           else
217             putint(Query(a,b));
218       }
219     return 0;
220 }