fork download
copy 
  1. #include "bits/stdc++.h"
  2. #include <ext/pb_ds/assoc_container.hpp>
  3. #include <ext/pb_ds/tree_policy.hpp>
  4.  
  5. using namespace std;
  6. using namespace __gnu_pbds;
  7.  
  8. #define ll long long
  9. #define ld double
  10. #define pb push_back
  11. #define pf push_front
  12. #define ppb pop_back
  13. #define ff first
  14. #define ss second
  15. #define ins insert
  16. #define vll vector <ll>
  17. #define vvll vector <vll>
  18. #define vbool vector <bool>
  19. #define pll pair <ll,ll>
  20. #define vpll vector <pll>
  21. #define YY cout<<"YES"
  22. #define NN cout<<"NO"
  23. #define yy cout<<"Yes"
  24. #define nn cout<<"No"
  25. #define set_bits __builtin_popcountll
  26. #define sz(v) (int)v.size()
  27. #define all(v) v.begin(), v.end()
  28. #define allr(v) v.rbegin(), v.rend()
  29. #define desc() greater <ll>()
  30. #define fill1(a,x) for (auto &it: a) it=x;
  31. #define fill2(a,x) for (auto &v: a) { for (auto &it: v) it=x; }
  32. #define endl '\n'   //not to be used in interactive problems
  33. #define random(l,r,T) uniform_int_distribution<T>(l,r)(rng)
  34. #define fastIO ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
  35.  
  36. mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
  37.  
  38. const int M = 1e9+7;
  39. const int MM = 998244353;
  40. const int N = 2e5+7;
  41. const ll inf = 1e18;
  42. const ld eps = 1e-9;
  43. #define PI 3.141592653589793238462
  44.  
  45. int dx[]={0,1,0,-1};
  46. int dy[]={1,0,-1,0};
  47. int Dx[]={0,1,1,1,0,-1,-1,-1};
  48. int Dy[]={1,1,0,-1,-1,-1,0,1};
  49.  
  50.  
  51. // ********     Policy Based Data Structures     ********
  52. template<class T> using oset = tree<T,null_type,less<T>,rb_tree_tag,tree_order_statistics_node_update>;
  53. template<class key, class value> using omap = tree <key,value,less<key>,rb_tree_tag,tree_order_statistics_node_update>;
  54. // find_by_order(k) -> returns iterator to kth element from start 0
  55. // order_of_key(k) -> returns count of elements < k
  56.  
  57.  
  58. // ********     Debug Helper     ********
  59. vector<string> vec_splitter(string s) {
  60.     s += ',';
  61.     vector<string> res;
  62.     while(!s.empty()) {
  63.         res.push_back(s.substr(0, s.find(',')));
  64.         s = s.substr(s.find(',') + 1);
  65.     }
  66.     return res;
  67. }
  68.  
  69. void debug_out(
  70. vector<string> __attribute__ ((unused)) args,
  71. __attribute__ ((unused)) int idx, 
  72. __attribute__ ((unused)) int LINE_NUM) { cerr << endl; } 
  73. template <typename Head, typename... Tail>
  74. void debug_out(vector<string> args, int idx, int LINE_NUM, Head H, Tail... T) {
  75.     if(idx > 0) cerr << ", "; else cerr << "Line(" << LINE_NUM << ") ";
  76.     stringstream ss; ss << H;
  77.     cerr << args[idx] << " = " << ss.str();
  78.     debug_out(args, idx + 1, LINE_NUM, T...);
  79. }
  80.  
  81. template<typename C, 
  82.          typename T = std::decay_t<decltype(*begin(std::declval<C>()))>,
  83.          typename std::enable_if<!std::is_same<C, std::string>::value>::type* = nullptr
  84.          >
  85. std::ostream &operator<<(std::ostream &os, const C &container)
  86. {
  87.   bool first = true;
  88.   std::stringstream ss; 
  89.   ss << '[';
  90.   for(const auto &x : container){
  91.     if (!first){
  92.       ss << ", ";
  93.     }
  94.     first = false;
  95.     ss << x;
  96.   }
  97.   ss << ']';
  98.   return os << ss.str();
  99. }
  100.  
  101. #ifndef ONLINE_JUDGE
  102.     #define debug(...) debug_out(vec_splitter(#__VA_ARGS__), 0, __LINE__, __VA_ARGS__)
  103. #else
  104.     #define debug(...) 42
  105. #endif
  106.  
  107.  
  108. // ********     Useful Fuctions     ********
  109. ll gcd(ll a, ll b) { while (b) {a %= b; swap(a,b);} return a; }
  110. ll lcm(ll a, ll b) { ll g=gcd(a,b); ll res=a*(b/g); return res; }
  111. ll extended_gcd(ll a, ll b, ll &x, ll &y) { if (b==0) { x=1; y=0; return a; } ll x1,y1; ll g=extended_gcd(b,a%b,x1,y1); x=y1; y=x1-y1*(a/b); return g; }
  112. ll binExp(ll a, ll b, ll m=M) { a = a % m; ll res = 1; while (b) { if (b&1) { res=(res * a) % m; } a=(a * a) % m; b>>=1; } return res; }
  113. ll mod_inv(ll a, ll m=M) { a = a % m; return binExp(a,m-2,m); }        // only for prime m
  114. ll mod_add(ll a, ll b, ll m=M) { a = a % m; b = b % m; return (((a + b) % m) + m) % m; }
  115. ll mod_sub(ll a, ll b, ll m=M) { a = a % m; b = b % m; return (((a - b) % m) + m) % m; }
  116. ll mod_mul(ll a, ll b, ll m=M) { a = a % m; b = b % m; return (((a * b) % m) + m) % m; }
  117. ll mod_div(ll a, ll b, ll m=M) { a = a % m; ll binv = mod_inv(b,m); return (((a * binv) % m) + m) % m; }
  118. ll sqrtll(ll n) { ll lo=0,hi=3037000499; while (hi-lo>1) { ll m=(hi+lo)/2; if (m*m<=n) { lo=m; } else { hi=m-1; }} if (hi*hi<=n) { return hi; } return lo; }
  119. ld sqrtld(ll n) { ld lo=0,hi=3037000499; while (hi-lo>eps) { ld m=(hi+lo)/2; if ((n-m*m)>eps) { lo=m; } else { hi=m-eps; }} return lo; }
  120. ll cbrtll(ll n) { ll lo=0,hi=2097151; while (hi-lo>1) { ll m=(hi+lo)/2; if (m*m*m<=n) { lo=m; } else { hi=m-1; }} if (hi*hi*hi<=n) { return hi; } return lo; }
  121. ld cbrtld(ll n) { ld lo=0,hi=2097151; while (hi-lo>eps) { ld m=(hi+lo)/2; if ((n-m*m*m)>eps) { lo=m; } else { hi=m-eps; }} return lo; }
  122. void init_usaco() { freopen("input.txt","r",stdin); freopen("output.txt","w",stdout); }
  123.  
  124.  
  125. // ********     Input/Output Template     *********
  126. template<typename T> istream& operator >>(istream &in,vector<T> &v){ for(auto &x:v) in>>x; return in;}
  127. template<typename T> ostream& operator <<(ostream &out,const vector<T> &v){ for(auto &x:v) out<<x<<' '; return out;}
  128. template<typename T1,typename T2> istream& operator >>(istream &in,pair<T1,T2> &p){ in>>p.first>>p.second; return in;}
  129. template<typename T1,typename T2> ostream& operator <<(ostream &out,const pair<T1,T2> &p){ out<<p.first<<' '<<p.second; return out;}
  130.  
  131.  
  132. // ********     ACCEPTED     ********
  133. vll color,val;
  134. vll parent,sbtreeSize,lvl,chain,chainHead,positionInHldArr;
  135. vll hldArr,hldValArr;
  136. vll inTime,outTime;
  137. vvll up;
  138. ll pos,chainId;
  139.  
  140. typedef struct Node
  141. {
  142.     ll w=0,b=0,sumW=0,sumB=0;
  143. }Node;
  144.  
  145. vector <Node> arr;               // { sum of white, sum of black}
  146. vll lazyarr;            // { 0 -> no update, 1-> update to black node 2-> update to white node}
  147.  
  148. void dfs(ll node, ll par, ll l, vll adj[])
  149. {
  150.     sbtreeSize[node]=1;
  151.     parent[node]=par;
  152.     up[node][0]=par;
  153.     lvl[node]=l;
  154.     // binary lifting
  155.     for (ll i=1;i<20;i++)
  156.     {
  157.         if (up[node][i-1]!=-1)
  158.             up[node][i]=up[up[node][i-1]][i-1];
  159.         else
  160.             up[node][i]=-1;
  161.     }
  162.     for (auto &it: adj[node])
  163.     {
  164.         if (it!=par)
  165.         {
  166.             dfs(it,node,l+1,adj);
  167.             sbtreeSize[node]+=sbtreeSize[it];
  168.         }
  169.     }
  170. }
  171.  
  172. ll lift_node(ll node, ll k)
  173. {
  174.     for (ll i=19;i>=0 && node!=-1;i--)
  175.     {
  176.         if (k&(1LL<<i))
  177.             node=up[node][i];
  178.     }
  179.     return node;
  180. }
  181.  
  182. ll LCA(ll a, ll b)
  183. {
  184.     if (lvl[a]<lvl[b]) swap(a,b);
  185.     a=lift_node(a,lvl[a]-lvl[b]);
  186.     if (a==b) return a;
  187.     for (ll i=19;i>=0;i--)
  188.     {
  189.         ll para=up[a][i];
  190.         ll parb=up[b][i];
  191.         if (para!=parb)
  192.         {
  193.             a=para;
  194.             b=parb;
  195.         }
  196.     }
  197.     return lift_node(a,1);
  198. }
  199.  
  200. void HLD(ll node, ll par, vll adj[])
  201. {
  202.     chain[node]=chainId;
  203.     hldArr[pos]=node;
  204.     hldValArr[pos]=val[node];
  205.     positionInHldArr[node]=pos;
  206.     inTime[node]=pos++;
  207.     ll heavyChild=-1,heavySize=0;
  208.     for (auto &it: adj[node])
  209.     {
  210.         if (it!=par)
  211.         {
  212.             if (sbtreeSize[it]>heavySize)
  213.             {
  214.                 heavySize=sbtreeSize[it];
  215.                 heavyChild=it;
  216.             }
  217.         }
  218.     }
  219.     if (heavySize>0) HLD(heavyChild,node,adj);  // not leaf, move to heavy edge
  220.     for (auto &it: adj[node])
  221.     {
  222.         if (it!=par && it!=heavyChild)
  223.         {
  224.             chainHead[++chainId]=it;
  225.             HLD(it,node,adj);
  226.         }
  227.     }
  228.     outTime[node]=pos-1;
  229. }
  230.  
  231. void merge(ll ind, ll left, ll right)
  232. {
  233.     arr[ind].w=arr[left].w+arr[right].w;
  234.     arr[ind].b=arr[left].b+arr[right].b;
  235.     arr[ind].sumW=arr[left].sumW+arr[right].sumW;
  236.     arr[ind].sumB=arr[left].sumB+arr[right].sumB;
  237. }
  238.  
  239. void update_lazy(ll ind, ll lo, ll hi)
  240. {
  241.     if (lazyarr[ind]==0) return;
  242.     ll upd=lazyarr[ind];
  243.     lazyarr[ind]=0;
  244.     if (upd==1)
  245.     {
  246.         arr[ind].w+=arr[ind].b;
  247.         arr[ind].sumW+=arr[ind].sumB;
  248.         arr[ind].b=arr[ind].sumB=0;
  249.     }
  250.     else
  251.     {
  252.         arr[ind].b+=arr[ind].w;
  253.         arr[ind].sumB+=arr[ind].sumW;
  254.         arr[ind].w=arr[ind].sumW=0;
  255.     }
  256.     if (lo!=hi) lazyarr[2*ind+1]=upd, lazyarr[2*ind+2]=upd;
  257. }
  258.  
  259. void buildSegTree(ll ind, ll lo, ll hi)
  260. {
  261.     if (lo==hi)
  262.     {
  263.         arr[ind].w=1;
  264.         arr[ind].sumW=hldValArr[lo];
  265.         arr[ind].b=arr[ind].sumB=0;
  266.         return;
  267.     }
  268.     ll mid=(hi+lo)/2;
  269.     buildSegTree(2*ind+1,lo,mid);
  270.     buildSegTree(2*ind+2,mid+1,hi);
  271.     merge(ind,2*ind+1,2*ind+2);
  272. }
  273.  
  274. /*Query:
  275. 1 -> get white count
  276. 2 -> get black count
  277. 3 -> get sum white
  278. 4 -> get sum black*/
  279.  
  280. ll querySegTree(ll ind, ll lo, ll hi, ll l, ll r, ll d)
  281. {
  282.     // check for update
  283.     update_lazy(ind,lo,hi);
  284.     if (lo>r || hi<l) return 0;
  285.     else if (lo>=l && hi<=r)
  286.     {
  287.         if (d==1) return arr[ind].w;
  288.         else if (d==2) return arr[ind].b;
  289.         else if (d==3) return arr[ind].sumW;
  290.         else return arr[ind].sumB;
  291.     }
  292.     ll mid=(hi+lo)/2;
  293.     ll left=querySegTree(2*ind+1,lo,mid,l,r,d);
  294.     ll right=querySegTree(2*ind+2,mid+1,hi,l,r,d);
  295.     ll ans=left+right;
  296.     return ans;
  297. }
  298.  
  299. /*Update:
  300. 1 -> force all white
  301. 2 -> force all black*/
  302.  
  303. void updateSegTree(ll ind, ll lo, ll hi, ll l, ll r, ll x)
  304. {
  305.     // check for update
  306.     update_lazy(ind,lo,hi);
  307.     if (lo>r || hi<l) return;
  308.     else if (lo>=l && hi<=r)
  309.     {
  310.         if (x==1)
  311.         {
  312.             arr[ind].w+=arr[ind].b;
  313.             arr[ind].sumW+=arr[ind].sumB;
  314.             arr[ind].b=arr[ind].sumB=0;
  315.         }
  316.         else
  317.         {
  318.             arr[ind].b+=arr[ind].w;
  319.             arr[ind].sumB+=arr[ind].sumW;
  320.             arr[ind].w=arr[ind].sumW=0;
  321.         }
  322.         if (lo!=hi) lazyarr[2*ind+1]=x, lazyarr[2*ind+2]=x;
  323.         return;
  324.     }
  325.     ll mid=(hi+lo)/2;
  326.     updateSegTree(2*ind+1,lo,mid,l,r,x);
  327.     updateSegTree(2*ind+2,mid+1,hi,l,r,x);
  328.     merge(ind,2*ind+1,2*ind+2);
  329. }
  330.  
  331. void hld_query1()
  332. {
  333.     ll n=sz(hldArr);
  334.     ll x,y; cin>>x>>y; x--; y--;
  335.     ll lca=LCA(x,y);
  336.     ll whiteCnt=0,blackCnt=0;
  337.     ll a=x,b=y;
  338.     while (chain[a]!=chain[lca])
  339.     {
  340.         whiteCnt+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[a]]],positionInHldArr[a],1);
  341.         blackCnt+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[a]]],positionInHldArr[a],2);
  342.         a=up[chainHead[chain[a]]][0];
  343.     }
  344.     whiteCnt+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[a],1);
  345.     blackCnt+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[a],2);
  346.     while (chain[b]!=chain[lca])
  347.     {
  348.         whiteCnt+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[b]]],positionInHldArr[b],1);
  349.         blackCnt+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[b]]],positionInHldArr[b],2);
  350.         b=up[chainHead[chain[b]]][0];
  351.     }
  352.     whiteCnt+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[b],1);
  353.     blackCnt+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[b],2);
  354.     if (querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[lca],1)==1) whiteCnt--;
  355.     else blackCnt--;
  356.     if (whiteCnt>blackCnt)
  357.     {
  358.         // force white in path
  359.         a=x,b=y;
  360.         while (chain[a]!=chain[lca])
  361.         {
  362.             updateSegTree(0,0,n-1,positionInHldArr[chainHead[chain[a]]],positionInHldArr[a],1);
  363.             a=up[chainHead[chain[a]]][0];
  364.         }
  365.         updateSegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[a],1);
  366.         while (chain[b]!=chain[lca])
  367.         {
  368.             updateSegTree(0,0,n-1,positionInHldArr[chainHead[chain[b]]],positionInHldArr[b],1);
  369.             b=up[chainHead[chain[b]]][0];
  370.         }
  371.         updateSegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[b],1);
  372.     }
  373.     else if (blackCnt>whiteCnt)
  374.     {
  375.         // force black in path
  376.         a=x,b=y;
  377.         while (chain[a]!=chain[lca])
  378.         {
  379.             updateSegTree(0,0,n-1,positionInHldArr[chainHead[chain[a]]],positionInHldArr[a],2);
  380.             a=up[chainHead[chain[a]]][0];
  381.         }
  382.         updateSegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[a],2);
  383.         while (chain[b]!=chain[lca])
  384.         {
  385.             updateSegTree(0,0,n-1,positionInHldArr[chainHead[chain[b]]],positionInHldArr[b],2);
  386.             b=up[chainHead[chain[b]]][0];
  387.         }
  388.         updateSegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[b],2);
  389.     }
  390. }
  391.  
  392. void hld_query2()
  393. {
  394.     ll n=sz(hldArr);
  395.     ll x; cin>>x; x--;
  396.     // color subtree x with black
  397.     updateSegTree(0,0,n-1,inTime[x],outTime[x],2);
  398. }
  399.  
  400. ll hld_query3()
  401. {
  402.     ll n=sz(hldArr);
  403.     ll x,y; cin>>x>>y; x--; y--;
  404.     ll lca=LCA(x,y);
  405.     ll whiteSum=0,blackSum=0;
  406.     ll a=x,b=y;
  407.     while (chain[a]!=chain[lca])
  408.     {
  409.         whiteSum+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[a]]],positionInHldArr[a],3);
  410.         blackSum+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[a]]],positionInHldArr[a],4);
  411.         a=up[chainHead[chain[a]]][0];
  412.     }
  413.     whiteSum+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[a],3);
  414.     blackSum+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[a],4);
  415.     while (chain[b]!=chain[lca])
  416.     {
  417.         whiteSum+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[b]]],positionInHldArr[b],3);
  418.         blackSum+=querySegTree(0,0,n-1,positionInHldArr[chainHead[chain[b]]],positionInHldArr[b],4);
  419.         b=up[chainHead[chain[b]]][0];
  420.     }
  421.     whiteSum+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[b],3);
  422.     blackSum+=querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[b],4);
  423.     if (querySegTree(0,0,n-1,positionInHldArr[lca],positionInHldArr[lca],1)==1) whiteSum-=val[lca];
  424.     else blackSum-=val[lca];
  425.     return max(whiteSum,blackSum);
  426. }
  427.  
  428. void solve()
  429. {
  430.     ll n,m; cin>>n>>m;
  431.     hldArr.resize(n);
  432.     hldValArr.resize(n);
  433.     vll adj[n];
  434.     for (ll i=0;i<n-1;i++)
  435.     {
  436.         ll u,v; cin>>u>>v; u--; v--;
  437.         adj[u].push_back(v);
  438.         adj[v].push_back(u);
  439.     }
  440.     val.resize(n); cin>>val;
  441.     color.resize(n,0);
  442.     parent.resize(n);
  443.     sbtreeSize.resize(n);
  444.     lvl.resize(n);
  445.     chain.resize(n);
  446.     chainHead.resize(n);
  447.     positionInHldArr.resize(n);
  448.     inTime.resize(n);
  449.     outTime.resize(n);
  450.     up.resize(n, vll (20,-1));
  451.     arr.resize(4*n);
  452.     lazyarr.resize(4*n,0);
  453.     dfs(0,-1,0,adj);
  454.     pos=chainId=0;
  455.     HLD(0,-1,adj);
  456.     debug(hldArr);
  457.     debug(hldValArr);
  458.     debug(inTime);
  459.     debug(outTime);
  460.     buildSegTree(0,0,n-1);
  461.     while (m--)
  462.     {
  463.         ll q; cin>>q;
  464.         if (q==1) hld_query1();
  465.         else if (q==2) hld_query2();
  466.         else cout<<hld_query3()<<endl;
  467.     }
  468. }
  469.  
  470. int main()
  471. {
  472.     #ifndef ONLINE_JUDGE
  473.         // freopen("input.txt","r",stdin);
  474.         // freopen("output.txt","w",stdout);
  475.         freopen("error.txt","w",stderr);
  476.         clock_t clk = clock();
  477.     #endif
  478.     // init_usaco();
  479.     fastIO;
  480.     int t=1;
  481.     // cin>>t;
  482.     for (int test=1;test<=t;test++)
  483.     {
  484.         // cout<<"Case #"<<test<<": ";
  485.         solve();
  486.         cout<<endl;
  487.     }
  488.     #ifndef ONLINE_JUDGE
  489.           cerr << '\n'<<"Time (in s): " << double(clock() - clk) * 1.0 / CLOCKS_PER_SEC << '\n';
  490.     #endif
  491.     return 0;
  492. }
  493.  
Runtime error #stdin #stdout 0.01s 5428KB
comments ()
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/SgC4ZB
language:
C++ (gcc 8.3)
created:
2 years ago
visibility:
secret

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!