2016shenyang-1002-HDU5893-List wants to travel-树链剖分+线段树维护不同区间段个数

  1 /*--------------------------------------------------------------------------------------*/
  2 
  3 #include <algorithm>
  4 #include <iostream>
  5 #include <cstring>
  6 #include <ctype.h>
  7 #include <cstdlib>
  8 #include <cstdio>
  9 #include <vector>
 10 #include <string>
 11 #include <queue>
 12 #include <stack>
 13 #include <cmath>
 14 #include <set>
 15 #include <map>
 16 
 17 //debug function for a N*M array
 18 #define debug_map(N,M,G) printf("\n");for(int i=0;i<(N);i++) 19 {for(int j=0;j<(M);j++){ 20 printf("%d",G[i][j]);}printf("\n");}
 21 //debug function for int,float,double,etc.
 22 #define debug_var(X) cout<<#X"="<<X<<endl;
 23 #define LL long long
 24 const int INF = 0x3f3f3f3f;
 25 const LL LLINF = 0x3f3f3f3f3f3f3f3f;
 26 /*--------------------------------------------------------------------------------------*/
 27 using namespace std;
 28 
 29 int N,M,T;
 30 const int maxn = 2e5+100;
 31 
 32 struct Edge{
 33     int x,y,val;
 34 }e[maxn];
 35 
 36 vector <int> G[maxn];
 37 
 38 int val[maxn],dep[maxn],siz[maxn],top[maxn],id[maxn],son[maxn],fa[maxn];
 39 int topw;
 40 
 41 void init()
 42 {
 43     for(int i=0;i<maxn;i++) G[i].clear();
 44     topw = 0;
 45 }
 46 
 47 void dfs_1(int u,int f,int d)
 48 {
 49     fa[u] = f;
 50     son[u] = 0;
 51     siz[u] = 1;
 52     dep[u] = d;
 53     for(int i=0;i<G[u].size();i++) if(G[u][i] != f)
 54     {
 55         dfs_1(G[u][i],u,d+1);
 56         siz[u] += siz[G[u][i]];
 57         if(siz[son[u]] < siz[G[u][i] ])
 58         {
 59             son[u] = G[u][i];
 60         }
 61     }
 62 }
 63 
 64 void dfs_2(int u,int tp)
 65 {
 66     top[u] = tp;
 67     id[u] = ++topw;
 68     if(son[u]) dfs_2(son[u],tp);
 69     for(int i=0;i<G[u].size();i++) if(G[u][i]!=fa[u] && G[u][i] != son[u] )
 70     {
 71         dfs_2(G[u][i],G[u][i]);
 72     }
 73 }
 74 
 75 void debug()
 76 {
 77     for(int i=1;i<=N;i++)
 78     {
 79         printf("%d siz:%d son:%d dep:%d fa:%d ",i,siz[i],son[i],dep[i],fa[i]);
 80         printf("top:%d id:%d\n",top[i],id[i]);
 81     }
 82     return ;
 83 }
 84 
 85 /*-----------------------------------*/
 86 #define lson(x) (x<<1)
 87 #define rson(x) (x<<1|1)
 88 
 89 struct SegmentTree{
 90     int l,r;
 91     int lazy;
 92     int num,lcost,rcost;
 93 }segtree[4*maxn];
 94 
 95 void pushUp(int x)
 96 {
 97     segtree[x].num = segtree[lson(x)].num + segtree[rson(x)].num - (segtree[lson(x)].rcost==segtree[rson(x)].lcost? 1 : 0);
 98     segtree[x].lcost = segtree[lson(x)].lcost;
 99     segtree[x].rcost = segtree[rson(x)].rcost;
100     //printf("x:%d [%d,%d] num:%d\n",x,segtree[x].l,segtree[x].r,segtree[x].num);
101     //printf("lcost:%d rcost:%d\n",segtree[x].lcost,segtree[x].rcost);
102     //printf("l->rcost:%d r->lcost:%d\n",segtree[lson(x)].rcost,segtree[rson(x)].lcost);
103 }
104 
105 void pushDown(int x)
106 {
107     if(segtree[x].lazy != -1)
108     {
109         segtree[lson(x)].lcost = segtree[lson(x)].rcost = segtree[x].lazy;
110         segtree[rson(x)].lcost = segtree[rson(x)].rcost = segtree[x].lazy;
111         segtree[lson(x)].num = segtree[rson(x)].num = 1;
112         segtree[lson(x)].lazy = segtree[rson(x)].lazy = segtree[x].lazy;
113         segtree[x].lazy = -1;
114     }
115 }
116 
117 void Build(int l,int r,int x)
118 {
119     segtree[x].l = l;
120     segtree[x].r = r;
121     segtree[x].lazy = -1;
122     if(l == r)
123     {
124         segtree[x].num = 1;
125         segtree[x].lcost = segtree[x].rcost = val[l];
126         return ;
127     }
128     int mid = (l+r)>>1;
129     Build(l,mid,lson(x));
130     Build(mid+1,r,rson(x));
131     pushUp(x);
132 }
133 
134 void update(int x,int L,int R,int cost)
135 {
136     if(segtree[x].l >= L && segtree[x].r <= R)
137     {
138         segtree[x].lcost = segtree[x].rcost = cost;
139         segtree[x].num = 1;
140         segtree[x].lazy = cost;
141         return ;
142     }
143     pushDown(x);
144     int mid = (segtree[x].l + segtree[x].r) >> 1;
145     if(L <= mid) update(lson(x),L,R,cost);
146     if(R >  mid) update(rson(x),L,R,cost);
147     pushUp(x);
148     return ;
149 }
150 
151 int query(int x,int L,int R)
152 {
153     if(segtree[x].l >= L && segtree[x].r <= R)
154     {
155         return segtree[x].num;
156     }
157     pushDown(x);
158     int mid = (segtree[x].l + segtree[x].r) >> 1;
159     int ans = 0;
160 
161     if(R <= mid) ans = query(lson(x),L,R);
162     else if(L > mid) ans = query(rson(x),L,R);
163     else ans = query(lson(x),L,R) + query(rson(x),L,R) - (segtree[lson(x)].rcost==segtree[rson(x)].lcost ? 1 : 0) ;
164     pushUp(x);
165     return ans;
166 }
167 
168 int query_edge(int x,int pos)
169 {
170     if(segtree[x].l == segtree[x].r)
171     {
172         return segtree[x].lcost;
173     }
174     pushDown(x);
175     int res = 0;
176     int mid = (segtree[x].l + segtree[x].r) >> 1;
177     if(pos <= mid) res =  query_edge(lson(x),pos);
178     else           res =  query_edge(rson(x),pos);
179     pushUp(x);
180     return res ;
181 }
182 
183 int Find(int u,int v)
184 {
185     int ans = 0,fu = top[u],fv = top[v];
186     int last_u=-1 , last_v = -1;
187     int last_u_color,last_v_color;
188     //printf("%d->%d\n",u,v);
189 
190     while(fu != fv)
191     {
192         if(dep[fu] < dep[fv])
193         {
194             swap(fu,fv);swap(u,v);
195             swap(last_u,last_v);
196             swap(last_u_color,last_v_color);
197         }
198 
199         //printf("query:[%d,%d] %d->%d ",id[fu],id[u],u,fu);
200         //printf("%d\n",query(1,id[fu],id[u]));
201         ans += query(1,id[fu],id[u]);
202         if(last_u == -1)
203         {
204             last_u = fu;
205             last_u_color = query_edge(1,id[fu]);
206         }
207         else
208         {
209             if(last_u != -1 && fa[last_u] == u)
210             {
211                 //printf("u sub:%d\n",(last_u_color==query_edge(1,id[u]) ? 1 : 0));
212                 ans -= (last_u_color==query_edge(1,id[u]) ? 1 : 0);
213                 last_u = fu;
214                 last_u_color = query_edge(1,id[fu]);
215             }
216         }
217         u = fa[fu];
218         fu = top[u];
219     }
220     if(u == v)
221     {
222         if(last_u != -1 && last_v != -1 && last_u_color==last_v_color) ans -= 1;
223         return ans;
224     }
225     if(dep[u] < dep[v])
226     {
227         swap(u,v);
228         swap(last_u,last_v);
229         swap(last_u_color,last_v_color);
230     }
231     //printf("query:[%d,%d] %d->%d ",id[son[v]],id[u],u,son[v]);
232     ans += query(1,id[son[v] ],id[u]);
233     //printf("%d*\n",query(1,id[son[v]],id[u]));
234 
235     //printf("last_u:%d last_v:%d\n",last_u,last_v);
236     if(last_u != -1 && fa[last_u] == u)
237     {
238         //printf("u sub:%d last_u_color:%d u_color:%d\n",(last_u_color==query_edge(1,id[u]) ? 1 : 0),last_u_color,query_edge(1,id[u]));
239         ans -= (last_u_color==query_edge(1,id[u]) ? 1 : 0);
240         last_u = fu;
241         last_u_color = query_edge(1,id[u]);
242     }
243 
244     if(last_v != -1 && fa[last_v] == v)
245     {
246         //printf("v sub:%d last_v_color:%d son[v]_color:%d\n",(last_v_color==query_edge(1,id[son[v]]) ? 1 : 0),last_v_color,query_edge(1,id[son[v]] ));
247         ans -= (last_v_color==query_edge(1,id[son[v]]) ? 1 : 0);
248         last_v = fv;
249         last_v_color = query_edge(1,id[son[v]]);
250     }
251 
252     return ans;
253 }
254 
255 void Change(int u,int v,int cost)
256 {
257     int fu = top[u],fv = top[v];
258     //printf("%d->%d\n",u,v);
259     while(fu != fv)
260     {
261         if(dep[fu] < dep[fv])
262         {
263             swap(fu,fv);swap(u,v);
264         }
265         //printf("change:[%d,%d] %d->%d %d\n",id[fu],id[u],u,fu,cost);
266         update(1,id[fu],id[u],cost);
267         u = fa[fu];
268         fu = top[u];
269     }
270     if(u == v) return ;
271     if(dep[u] < dep[v]) swap(u,v);
272     //printf("change:[%d,%d] %d->%d %d\n",id[son[v]],id[u],u,son[u],cost);
273     update(1,id[son[v]],id[u],cost);
274     return ;
275 }
276 
277 int main()
278 {
279     //freopen("input","r",stdin);
280     while(~scanf("%d%d",&N,&M))
281     {
282         init();
283         int u,v,w;
284         for(int i=1;i<N;i++)
285         {
286             scanf("%d%d%d",&u,&v,&w);
287             e[i].x = u;e[i].y = v;e[i].val = w;
288             G[u].push_back(v);
289             G[v].push_back(u);
290         }
291         topw = 0;
292         dfs_1(1,0,1);
293         dfs_2(1,1);
294         //debug();
295         for(int i=1;i<N;i++)
296         {
297             if(dep[e[i].x] < dep[e[i].y]) swap(e[i].x,e[i].y);
298             val[id[e[i].x]] = e[i].val;
299         }
300 
301         Build(1,topw,1);
302         char op[20];
303         //printf("-----------------\n");
304         for(int i=0;i<M;i++)
305         {
306             scanf("%s",op);
307             if(op[0] == ‘Q‘)
308             {
309                 scanf("%d%d",&u,&v);
310                 printf("%d\n",Find(u,v));
311             }
312             else
313             {
314                 scanf("%d%d%d",&u,&v,&w);
315                 Change(u,v,w);
316             }
317             //puts("");
318         }
319     }
320 }