ec1117

Untitled

Dec 13th, 2020
199
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.27 KB | None | 0 0
  1. template<int SZ> struct centroidDecomp {
  2.     int N;
  3.     bool done[SZ];
  4.     int sub[SZ], par[SZ], cen[SZ];
  5.     ll ans[SZ];
  6.     ll dist[SZ][20], tmp[SZ];
  7.     vpi adj[SZ];
  8.    
  9.     // INITIALIZE
  10.    
  11.     void addEdge(int a, int b, int d) { adj[a].pb({b,d}), adj[b].pb({a,d}); }
  12.    
  13.     void dfs (int no) {
  14.         sub[no] = 1;
  15.         for (pi i: adj[no]) if (!done[i.f] && i.f != par[no]) {
  16.             par[i.f] = no;
  17.             dfs(i.f);
  18.             sub[no] += sub[i.f];
  19.         }
  20.     }
  21.    
  22.     void genDist(int par, int no, int t, ll dis) {
  23.         dist[no][tmp[no]++] = dis;
  24.         for (pi i: adj[no]) if (!done[i.f] && i.f != par) {
  25.             cen[i.f] = t;
  26.             genDist(no,i.f,t,dis+i.s);
  27.         }
  28.     }
  29.    
  30.     int getCentroid(int x) {
  31.         par[x] = -1; dfs(x);
  32.         int sz = sub[x];
  33.         while (1) {
  34.             pi mx = {0,0};
  35.             for (pi i: adj[x]) if (!done[i.f] && i.f != par[x]) mx = max(mx,{sub[i.f],i.f});
  36.             if (mx.f*2 > sz) x = mx.s;
  37.             else return x;
  38.         }
  39.     }
  40.    
  41.     void solve (int x) {
  42.         x = getCentroid(x); done[x] = 1;
  43.         genDist(-1,x,x,0);
  44.         for (pi i: adj[x]) if (!done[i.f]) solve(i.f);
  45.     }
  46.    
  47.     void init() {
  48.         F0R(i,N) cen[i] = -1, ans[i] = INF;
  49.         solve(0);
  50.     }
  51.    
  52.     // QUERY
  53.    
  54.     void upd(int v, int x = 1) {
  55.         for (int V = v, ind = tmp[v]-1; V >= 0; V = cen[V], ind --) {
  56.             // cout << "OH " << v << " " << V << " " << x << "\n";
  57.             if (x == 1) ans[V] = min(ans[V],dist[v][ind]);
  58.             else ans[V] = INF;
  59.         }
  60.     }
  61.    
  62.     ll query(int v) {
  63.         ll ret = INF;
  64.         for (int V = v, ind = tmp[v]-1; V >= 0; V = cen[V], ind --) {
  65.             // cout << "AH " << V << " " << v << "\n";
  66.             ret = min(ret,ans[V]+dist[v][ind]);
  67.         }
  68.         return ret;
  69.     }
  70. };
  71.  
  72. centroidDecomp<MX> C;
  73.  
  74. void Init(int N, int A[], int B[], int D[]) {
  75.     C.N = N;
  76.     F0R(i,N-1) C.addEdge(A[i],B[i],D[i]);
  77.     C.init();
  78. }
  79.  
  80. long long Query(int S, int X[], int T, int Y[]) {
  81.     ll ans = INF;
  82.     if (S > T) {
  83.         swap(S,T);
  84.         swap(X,Y);
  85.     }
  86.     F0R(i,S) C.upd(X[i]);
  87.     F0R(i,T) ans = min(ans,C.query(Y[i]));
  88.     F0R(i,S) C.upd(X[i],-1);
  89.     return ans;
  90. }
  91.  
Advertisement
Add Comment
Please, Sign In to add comment