(13511) 트리와 쿼리 2

https://www.acmicpc.net/problem/13511

단계별로 풀어보기

47단계(최소 공통 조상) 5번째

 

 

 

경로의 비용은 이전 문제에서 최대값, 최소값을 저장한 것처럼 구간의 경로 비용을 저장하고 희소 테이블을 이용해 점프할 때마다 값을 더한 가중치를 반환하면 된다.

이제 경로상 k번째 정점을 찾는 쿼리를 해결해야 하는데 LCA를 구했다면 깊이 배열을 이용해 u에서 LCA까지 지나는 노드 수를 알 수 있다. 이 값이 k보다 크거나 같다면 u에서 k-1만큼(1번째 정점은 u이므로) 조상 노드로 점프한 결과를, k보다 크다면 반대로 v에서 u로 갈 때 (구간 내 총 정점 수-k)번째로 지나는 정점을 구하는 것과 같으므로 v에서 같은 방식으로 결과를 반환하면 된다.

 

 

 

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
using namespace std;

int main(){
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  
  int n, m;
  
  cin >> n;
  
  int h= (int)floor(log2(n))+1;
  vector adj(n+1, vector<pair<int,int>>());
  vector st(n+1, vector<pair<int, long>> (h));
  vector<int> depth(n+1);
  
  for(int i=1;i<n;i++){
    int u, v, w;
    
    cin >> u >> v >> w;
    
    adj[u].push_back({w, v});
    adj[v].push_back({w, u});
  }
  
  auto dfs= [&](auto self, int now, int prev) -> void {
    depth[now]= depth[prev]+1;
    st[now][0].first= prev;
    
    for(int i=1;i<h;i++){
      st[now][i].first= st[st[now][i-1].first][i-1].first;
      st[now][i].second= st[now][i-1].second+st[st[now][i-1].first][i-1].second;
    }
    
    for(auto next:adj[now]){
      if(next.second==prev)  continue;
      
      st[next.second][0].second= next.first;
      self(self, next.second, now);
    }
  };
  
  dfs(dfs, 1, 0);
  
  auto lca= [&](int u, int v, int k= -1){
    int x= u, y= v;
    long dist= 0;
    
    if(depth[x]<depth[y])  swap(x, y);
    
    for(int i=h-1;i>=0;i--){
      if(depth[st[x][i].first]>=depth[y]){
        dist+= st[x][i].second;
        x= st[x][i].first;
      }
    }
    
    if(x!=y){
      for(int i=h-1;i>=0;i--){
        if(st[x][i].first!=st[y][i].first){
          dist+= st[x][i].second+st[y][i].second;
          x= st[x][i].first;
          y= st[y][i].first;
        }
      }
      
      dist+= st[x][0].second+st[y][0].second;
      x= st[x][0].first;
    }
    
    if(k<0)  return dist;
    
    if(depth[u]-depth[x]>=k){
      for(int i=h-1;i>=0;i--){
        if(k&(1<<i))  u= st[u][i].first;
      }
      
      return (long)u;
    }else{
      k= depth[u]+depth[v]-2*depth[x]-k;
      
      for(int i=h-1;i>=0;i--){
        if(k&(1<<i))  v= st[v][i].first;
      }
      
      return (long)v;
    }
  };
  
  cin >> m;
  
  while(m--){
    int qtype;
    
    cin >> qtype;
    
    if(qtype==1){
      int u, v;
      
      cin >> u >> v;
      
      cout << lca(u, v) << "\n";
    }else{
      int u, v, k;
      
      cin >> u >> v >> k;
      
      cout << lca(u, v, k-1) << "\n";
    }
  }
  
  return 0;
}