Advertisement
welleyth

A. Colors

Mar 18th, 2022
970
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 5.81 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2. #include <ext/pb_ds/assoc_container.hpp>
  3. #include <ext/pb_ds/tree_policy.hpp>
  4. using namespace __gnu_pbds;
  5. using namespace std;
  6. #define int long long
  7. #define pb push_back
  8. #define mp make_pair
  9.  
  10. #pragma GCC optimize("Ofast")
  11. #pragma GCC optimize("unroll-loops")
  12. #pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,tune=native")
  13. #pragma GCC target("avx,avx2")
  14.  
  15. constexpr int INF = (int)1e9;
  16. constexpr int N = (int)2e5 + 11;
  17. constexpr int MAXN = (int)3000 + 11;
  18. constexpr int md = (int)1e9+7;
  19.  
  20. void add(int& a,int b){
  21.     a += b;
  22.     if(a >= md) a -= md;
  23. }
  24.  
  25. vector<int> g[N];
  26. int sz[N];
  27. int parent[N];
  28. int depth[N];
  29. int heavy[N];
  30. int up[20][N];
  31. int tin[N],tout[N];
  32. int timer = 0;
  33. int col[N];
  34.  
  35. void dfs(int v,int pr = 0){
  36.     sz[v] = 1;
  37.     parent[v] = pr;
  38.     tin[v] = timer++;
  39.     up[0][v] = pr;
  40.     for(int i = 1; i < 20; i++){
  41.         up[i][v] = up[i-1][up[i-1][v]];
  42.     }
  43.     for(auto& x : g[v]){
  44.         if(x == pr)
  45.             continue;
  46.         depth[x] = depth[v] + 1;
  47.         dfs(x,v);
  48.         if(heavy[v] == 0 || sz[x] > sz[heavy[v]])
  49.             heavy[v] = x;
  50.         sz[v] += sz[x];
  51.     }
  52.     tout[v] = timer++;
  53.     return;
  54. }
  55. bool upper(int a,int b){
  56.     return tin[a] <= tin[b] && tout[a] >= tout[b];
  57. }
  58. int lca(int a,int b){
  59.     if(upper(a,b))
  60.         return a;
  61.     if(upper(b,a))
  62.         return b;
  63.     for(int i = 19; i >= 0; i--){
  64.         if(up[i][a] != 0 && !upper(up[i][a],b))
  65.             a = up[i][a];
  66.     }
  67.     return up[0][a];
  68. }
  69.  
  70. vector<int> a;
  71. int pos[N];
  72. int head[N];
  73.  
  74. void hld(int v,int h){
  75.     pos[v] = a.size();
  76.     a.push_back(v);
  77.     head[v] = h;
  78.     if(heavy[v] != 0)
  79.         hld(heavy[v],h);
  80.     for(auto& x : g[v]){
  81.         if(x == parent[v] || x == heavy[v])
  82.             continue;
  83.         hld(x,x);
  84.     }
  85.     return;
  86. }
  87.  
  88. struct node{
  89.     int ans = 0;
  90.     int col_r = 0,col_l = 0;
  91.     node(){}
  92.     node(int x){
  93.         ans = 0;
  94.         col_l = col_r = x;
  95.     }
  96. } t[4*N];
  97. int prom[4*N];
  98. node combine(node a,node b){
  99.     if(a.col_l == 0)
  100.         return b;
  101.     if(b.col_l == 0)
  102.         return a;
  103.     node c;
  104.     c.ans = a.ans + b.ans + (a.col_r != b.col_l);
  105.     c.col_l = a.col_l;
  106.     c.col_r = b.col_r;
  107.     return c;
  108. }
  109. void push(int v){
  110.     if(prom[v] == 0)
  111.         return;
  112.     t[v<<1] = node(prom[v]);
  113.     t[v<<1|1] = node(prom[v]);
  114.     prom[v<<1] = prom[v<<1|1] = prom[v];
  115.     prom[v] = 0;
  116.     return;
  117. }
  118.  
  119. void build(int v,int l,int r){
  120.     if(l == r){
  121.         t[v] = node(col[a[l]]);
  122.         return;
  123.     }
  124.     int m = (l+r)>>1;
  125.     build(v<<1,l,m);
  126.     build(v<<1|1,m+1,r);
  127.     t[v] = combine(t[v<<1],t[v<<1|1]);
  128. }
  129. void upd(int v,int l,int r,int tl,int tr,int col){
  130.     if(l > r || tl > tr)
  131.         return;
  132. //    cerr << "upd = " << v << " " << l+1 << " " << r+1 << " " << tl+1 << " " << tr+1 << " " << col << "\n";
  133.     if(l == tl && r == tr){
  134.         t[v] = node(col);
  135.         prom[v] = col;
  136.         return;
  137.     }
  138.     push(v);
  139.     int m = (l+r)>>1;
  140.     upd(v<<1,l,m,tl,min(tr,m),col);
  141.     upd(v<<1|1,m+1,r,max(tl,m+1),tr,col);
  142.     t[v] = combine(t[v<<1],t[v<<1|1]);
  143.     return;
  144. }
  145. node get(int v,int l,int r,int tl,int tr){
  146.     if(l > r || tl > tr)
  147.         return node();
  148. //    cerr << l+1 << " " << r+1 << ", ans = " << t[v].ans << "\n";
  149. //    cerr << "colors = " << t[v].col_l << " " << t[v].col_r << "\n";
  150.     if(l == tl && r == tr)
  151.         return t[v];
  152.     push(v);
  153.     int m = (l+r)>>1;
  154.     return combine(get(v<<1,l,m,tl,min(tr,m)),get(v<<1|1,m+1,r,max(tl,m+1),tr));
  155. }
  156. int n;
  157.  
  158. node query(int a,int b){ /// a is upper than b
  159.     if(upper(b,a))
  160.         swap(a,b);
  161.     node cur = node();
  162.     for(; head[a] != head[b]; b = parent[head[b]]){
  163.         cur = combine(get(1,0,n-1,pos[head[b]],pos[b]),cur);
  164.     }
  165.     cur = combine(get(1,0,n-1,pos[a],pos[b]),cur);
  166.     return cur;
  167. }
  168.  
  169. int get(int a,int b){
  170.     int L = lca(a,b);
  171.     node A = query(L,a);
  172.     node B = query(L,b);
  173. //    cerr << "get " << L << " " << a << " = " << A.ans << "\n";
  174. //    cerr << "get " << L << " " << b << " = " << B.ans << "\n\n";
  175. //    swap(B.col_l,B.col_r);
  176.     return A.ans + B.ans;
  177. }
  178. void upd(int a,int b,int c){ /// a is upper than b
  179.     if(upper(b,a))
  180.         swap(a,b);
  181.     for(; head[a] != head[b]; b = parent[head[b]]){
  182. //        if (depth[head[a]] > depth[head[b]])
  183. //            swap(a, b);
  184.         upd(1,0,n-1,pos[head[b]],pos[b],c);
  185.     }
  186.     upd(1,0,n-1,pos[a],pos[b],c);
  187.     return;
  188. }
  189. void updPath(int a,int b,int c){
  190.     int l = lca(a,b);
  191.     upd(l,a,c);
  192.     upd(l,b,c);
  193.     return;
  194. }
  195.  
  196. void solve(){
  197.     cin >> n;
  198.  
  199.     for(int i = 1; i < n; i++){
  200.         int a,b;
  201.         cin >> a >> b;
  202.         g[a].pb(b);
  203.         g[b].pb(a);
  204.     }
  205.  
  206.     dfs(1);
  207.     hld(1,1);
  208.     for(int i = 1; i <= n; i++){
  209.         cin >> col[i];
  210.     }
  211.     build(1,0,n-1);
  212. //    cerr << get(1,1) << "\n";
  213. //    return;
  214. //    for(auto& x : a){
  215. //        cerr << x << " ";
  216. //    }
  217. //    cerr << "\n";
  218. //    updPath(6,4,2);
  219. //    cerr << get(2,3) << "\n";
  220. //    return;
  221.     int q;
  222.     cin >> q;
  223.     while(q--){
  224.         int tp;
  225.         cin >> tp;
  226.         if(tp == 1){
  227.             int u,v,col;
  228.             cin >> u >> v >> col;
  229.             updPath(u,v,col);
  230.         }
  231.         if(tp == 2){
  232.             int u,c;
  233.             cin >> u >> c;
  234.             upd(1,0,n-1,pos[u],pos[u]+sz[u]-1,c);
  235.         }
  236.         if(tp == 3){
  237.             int u,v;
  238.             cin >> u >> v;
  239.             cout << get(u,v) << "\n";
  240.         }
  241.     }
  242.  
  243.     return;
  244. }
  245. signed main(){
  246. //    ifstream cin("colors.01.in");
  247. //    ofstream cout("output.txt");
  248.     ios::sync_with_stdio(false);cin.tie(nullptr);cout.tie(nullptr);
  249. //    freopen("colors.01.in","r",stdin);
  250. //    freopen("output.txt","w",stdout);
  251.     int tests = 1;
  252. //      cin >> tests;
  253.     while(tests--){
  254.         solve();
  255.     }
  256. }
  257. /**
  258. 1 2 3 4 5 6
  259.  
  260. **/
  261.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement