jw910731

Heavy Light Composition

Mar 4th, 2020
104
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 3.34 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2. #ifdef TEST
  3.     #define debug(...) printf(__VA_ARGS__)
  4. #else
  5.     #define debug(...) (void)0
  6. #endif
  7. #define EB emplace_back
  8. #define MP make_pair
  9. #define X first
  10. #define Y second
  11. using namespace std;
  12. using ll = long long;
  13. using llu = long long unsigned;
  14. using pii = pair<int,int>;
  15. /************************/
  16. #define MXV 100000
  17. #define LOW(x) (x&(-(x)))
  18. struct BIT{
  19.     vector<ll> orig;
  20.     vector<ll> arr;
  21.     BIT(){
  22.         orig.EB(0); // placeholder
  23.     }
  24.     void construct(){
  25.         arr.resize(orig.size());
  26.         for(int i = 1 ; i < int(orig.size()) ; ++i){
  27.             add(i, orig[i]);
  28.         }
  29.     }
  30.     void add(int i, int v){
  31.         for(;i<=int(size());i+=LOW(i)) arr[i] += v;
  32.     }
  33.     ll sum(int i){
  34.         ll sum = 0;
  35.         for(;i>0;i-=LOW(i)) sum += arr[i];
  36.         return sum;
  37.     }
  38.     // [0] is not used
  39.     size_t size(){return arr.size()-1;}
  40.     ll sum_all(){return sum(size());}
  41. };
  42. // graph
  43. map<int, int> graph[MXV+5]; // dest & weight
  44. map<int, BIT> bit4tree;
  45. pair<int, int> edges[MXV+5];
  46.  
  47. // Heavy Light Decomposition
  48. int size[MXV+5], maxCh[MXV+5], depth[MXV+5], pa[MXV+5];
  49. int chain_top[MXV+5], ord[MXV+5], chain_indx[MXV+5];
  50. int ord_top = 0;
  51. // HLD
  52. void preDfs(int n, int _pa){
  53.     size[n] = 1;
  54.     for(auto &ch : graph[n]){
  55.         if(ch.X != _pa){
  56.             pa[ch.X] = n;
  57.             depth[ch.X] = depth[n] + 1;
  58.             preDfs(ch.X, n);
  59.             if(maxCh[n] < 0 || size[maxCh[n]] < size[ch.X]) maxCh[n] = ch.X;
  60.             size[n] += size[ch.X];
  61.         }
  62.     }
  63. }
  64. // n=>current node id, top => chain top node id (glob), indx => in chain (local) id
  65. void decomp(int n, int top, int indx){
  66.     ord[n] = ord_top++;
  67.     chain_top[n] = top;
  68.     chain_indx[n] = indx;
  69.     if(maxCh[n] < 0){
  70.         return;
  71.     }
  72.     bit4tree[top].orig.EB(graph[n][maxCh[n]]);
  73.     decomp(maxCh[n], top, indx+1);
  74.     for(auto ch : graph[n]){
  75.         if(ch.X != maxCh[n] && ch.X != pa[n]){ // no back edge or max child edge
  76.             decomp(ch.X, ch.X, 0);
  77.         }
  78.     }
  79. }
  80. ll LCA(int u, int v){
  81.     ll sum = 0;
  82.     int tu = chain_top[u], tv = chain_top[v];
  83.     while(tu != tv){
  84.         if(depth[tu] < depth[tv]){
  85.             swap(tu, tv);
  86.             swap(u, v);
  87.         }
  88.         sum += bit4tree[tu].sum(chain_indx[u]) + graph[tu][pa[tu]];
  89.         debug("@@%d %d\n", bit4tree[tu].sum(chain_indx[u]), graph[tu][pa[tu]]);
  90.         tu = chain_top[u = pa[tu]];
  91.     }
  92.     if(depth[u] < depth[v]) swap(u, v); // assert u deeper
  93.     sum += bit4tree[tu].sum(chain_indx[u]) - bit4tree[tv].sum(chain_indx[v]);
  94.     debug("@@%d %d\n", bit4tree[tu].sum(chain_indx[u]), bit4tree[tv].sum(chain_indx[v]));
  95.     return sum;
  96. }
  97.  
  98. void init(){
  99.     memset(maxCh, 0xff, sizeof(maxCh));
  100.     memset(depth, 0x00, sizeof(depth));
  101.     memset(chain_top, 0xff, sizeof(chain_top));
  102.     memset(ord, 0x00, sizeof(ord));
  103. }
  104. int main(){
  105.     init();
  106.     int N, cnt = 0;
  107.     scanf("%d", &N);
  108.     for(int i = 1 ; i < N ; ++i){
  109.         int a, b, w;
  110.         scanf("%d%d%d", &a, &b, &w);
  111.         graph[a][b] = w;
  112.         graph[b][a] = w;
  113.         edges[cnt++] = make_pair(a, b);
  114.     }
  115.     preDfs(1, -1);
  116.     decomp(1, 1, 0);
  117.     for(auto &v : bit4tree){
  118.         v.Y.construct();
  119.     }
  120.     int Q;
  121.     scanf("%d", &Q);
  122.     for(int i = 0 ; i < Q ; ++i){
  123.         int op, a, b;
  124.         scanf("%d%d%d", &op, &a, &b);
  125.         if(op == 1){
  126.             int u, v;
  127.             tie(u, v) = edges[a-1];
  128.             int tmp = graph[u][v];
  129.             graph[u][v] = b;
  130.             graph[v][u] = b;
  131.             if(chain_top[u] == chain_top[v]){
  132.                 if(depth[u] < depth[v]){
  133.                     swap(u, v);
  134.                 }
  135.                 bit4tree[chain_top[u]].add(chain_indx[u], b-tmp);
  136.             }
  137.         }
  138.         if(op == 2){
  139.             ll ans = LCA(a, b);
  140.             printf("%lld\n", ans);
  141.         }
  142.     }
  143.     return 0;
  144. }
Add Comment
Please, Sign In to add comment