fork download
copy 
  1. #include<bits/stdc++.h>
  2. #define fast_az_fuk ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
  3. #define ll long long 
  4. #define ull unsigned ll
  5. #define ld long double 
  6. #define pb push_back 
  7. #define pf push_front
  8. #define dll deque<ll> 
  9. #define vll vector<ll>
  10. #define vvll vector<vll> 
  11. #define pll pair<ll,ll> 
  12. #define vpll vector<pll>
  13. #define dpll deque<pll>
  14. #define mapll map<ll,ll>
  15. #define endl "\n" 
  16. #define all(v) v.begin(),v.end() 
  17. #define ms(a,x) memset(a,x,sizeof(a))
  18. using namespace std;
  19. ll n;
  20. vvll adj; vvll up; vll tin,tout; ll timer; ll l;
  21. vll level; vvll weight; vll distFromRoot;
  22. class tupl{
  23. public:
  24.     ll num,lev,par;
  25.     tupl(ll a,ll b,ll c){
  26.         num = a; lev=b; par = c;
  27.     }
  28. };
  29. void dfs(ll node = 1,ll parent = 1,ll dist=0){
  30.     tin[node]=++timer;
  31.     up[node][0] = parent;
  32.     distFromRoot[node] = dist;
  33.     for(ll i=1;i<=l;i++){
  34.         up[node][i] = up[up[node][i-1]][i-1];
  35.     }
  36.     for(ll i=0;i<adj[node].size();i++){
  37.         ll x = adj[node][i];
  38.         if(x==parent) continue;
  39.         dfs(x,node,dist+weight[node][i]);
  40.     }
  41.     tout[node]=++timer;
  42. }
  43. void bfs(){
  44.     queue<tupl> q;
  45.     q.push(tupl(1,0,1));
  46.     while(!q.empty()){
  47.         tupl t = q.front(); q.pop();
  48.         level[t.num]=t.lev;
  49.         for(ll x:adj[t.num]){
  50.             if(x==t.par) continue;
  51.             q.push(tupl(x,t.lev+1,t.num));
  52.         }
  53.     }
  54. }
  55. bool isAncestor(ll u,ll v){
  56.     return (tin[u]<=tin[v]&&tout[u]>=tout[v]);
  57. }
  58. ll LCA(ll u,ll v)
  59. {
  60.     if(isAncestor(u,v)) return u;
  61.     else if(isAncestor(v,u)) return v;
  62.     for(ll i=l;i>=0;i--){
  63.         if(!isAncestor(up[u][i],v)){
  64.             u = up[u][i];
  65.         }
  66.     }
  67.     return up[u][0];
  68. }
  69. ll find(ll u,ll k){ // find the node, K distances up from u
  70.     for(ll i=l;i>=0;i--){
  71.         ll dist = (1LL<<i);
  72.         if(dist<=k){ u = up[u][i]; k-=dist; }
  73.     }
  74.     return u;
  75. }
  76. int32_t main()
  77. {
  78.     fast_az_fuk
  79.     ll testcase; cin>>testcase;
  80.     for(ll test=1;test<=testcase;test++)
  81.     {//cout<<"Case #"<<test<<": ";
  82.         cin>>n; adj.clear(); adj.resize(n+1); weight.clear(); weight.resize(n+1);
  83.         l = ceil(log2(n));
  84.         up.clear(); up.resize(n+1,vll(l+1));
  85.         tin.resize(n+1); tout.resize(n+1); level.resize(n+1); distFromRoot.resize(n+1); timer=0;
  86.         for(ll i=1;i<n;i++){
  87.             ll a,b; cin>>a>>b; ll w; cin>>w;
  88.             adj[a].pb(b); adj[b].pb(a); weight[a].pb(w); weight[b].pb(w);
  89.         }
  90.         dfs(); bfs();
  91.         while(1){
  92.             string s; cin>>s; if(s=="DONE") break; ll a,b; cin>>a>>b; ll k=-1;
  93.             if(s=="KTH") cin>>k;
  94.             if(k==-1){
  95.                 ll lc = LCA(a,b);
  96.                 cout<<(distFromRoot[a]+distFromRoot[b])-2*distFromRoot[lc]<<endl;
  97.             }
  98.             else{
  99.                 ll lc = LCA(a,b);
  100.                 ll numOfNodes = (level[a]+level[b]) - 2*level[lc]+1;
  101.                 if(k<=level[lc]-level[a]+1) cout<<find(a,k-1)<<endl;
  102.                 else cout<<find(b,numOfNodes-k)<<endl;
  103.             }
  104.         }
  105.         cout<<endl;
  106.     }
  107.     return 0;
  108. }
Success #stdin #stdout 0.01s 4876KB
comments (?)
stdin
copy 
2
6
1 2 1
2 4 1
2 5 2
1 3 1
3 6 2
DIST 4 6
KTH 4 6 4
DONE

3
1 2 1
2 3 1
DIST 1 3
KTH 1 3 2
DONE
stdout
copy 
5
3

2
2
https://ideone.com/NyVibL
language:
C++14 (gcc 8.3)
created:
3 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!