mickypinata

IPST59: Star Inter

Jul 30th, 2021
852
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. typedef long long lli;
  5. typedef pair<int, int> pii;
  6.  
  7. const int N = 2e5;
  8.  
  9. vector<pii> adj[N + 1];
  10. vector<int> path[N + 1], allPaths;
  11. int limDist;
  12.  
  13. void DFS(int u, int p, int i, int dist){
  14.     allPaths.push_back(dist);
  15.     path[i].push_back(dist);
  16.     for(pii nxt : adj[u]){
  17.         int v = nxt.first;
  18.         int w = nxt.second;
  19.         if(v != p){
  20.             DFS(v, u, i, dist + w);
  21.         }
  22.     }
  23. }
  24.  
  25. lli cntInside(int st){
  26.     lli cnt = 0;
  27.     int j = upper_bound(path[st].begin(), path[st].end(), path[st][0] + limDist) - path[st].begin();
  28.     for(int i = 0; i < path[st].size(); ++i){
  29.         while(j < path[st].size() && path[st][j] <= path[st][i] + limDist){
  30.             ++j;
  31.         }
  32.         cnt += j - i - 1;
  33.     }
  34.     return cnt;
  35. }
  36.  
  37. lli cntBetween(int st){
  38.     lli cnt = 0;
  39.     int j = upper_bound(path[st].begin(), path[st].end(), limDist - path[st][0]) - path[st].begin() - 1;
  40.     int k = upper_bound(allPaths.begin(), allPaths.end(), limDist - path[st][0]) - allPaths.begin() - 1;
  41.     for(int i = 0; i < path[st].size(); ++i){
  42.         while(j >= 0 && path[st][j] > limDist - path[st][i]){
  43.             --j;
  44.         }
  45.         while(k >= 0 && allPaths[k] > limDist - path[st][i]){
  46.             --k;
  47.         }
  48.         cnt += k - j;
  49.     }
  50.     return cnt;
  51. }
  52.  
  53. int main(){
  54.  
  55.     int nVertex;
  56.     scanf("%d%d", &nVertex, &limDist);
  57.     for(int i = 1; i < nVertex; ++i){
  58.         int u, v, w;
  59.         scanf("%d%d%d", &u, &v, &w);
  60.         adj[u].emplace_back(v, w);
  61.         adj[v].emplace_back(u, w);
  62.     }
  63.     int root, mx = 0;
  64.     for(int i = 1; i <= nVertex; ++i){
  65.         if(adj[i].size() > mx){
  66.             mx = adj[i].size();
  67.             root = i;
  68.         }
  69.     }
  70.     lli cnt = 0;
  71.     for(int i = 0; i < adj[root].size(); ++i){
  72.         int u = adj[root][i].first;
  73.         int w = adj[root][i].second;
  74.         DFS(u, root, i, w);
  75.         // To Center
  76.         int idx = upper_bound(path[i].begin(), path[i].end(), limDist) - path[i].begin();
  77.         cnt += idx;
  78.     }
  79.     sort(allPaths.begin(), allPaths.end());
  80.  
  81.     lli tmp = 0;
  82.     int sz = adj[root].size();
  83.     for(int i = 0; i < sz; ++i){
  84.         // Inside Paths
  85.         cnt += cntInside(i);
  86.         // Between Paths
  87.         tmp += cntBetween(i);
  88.     }
  89.     cnt += tmp / 2;
  90.     cout << cnt;
  91.  
  92.     return 0;
  93. }
  94.  
RAW Paste Data