SP913 QTREE2 - Query on a tree II

嘟嘟嘟


LCA水题，第二问看一下\(x\)到\(lca\)的路径长度是否够\(k - 1\)，不过的话就从\(y\)出发往上跳。

#include<cstdio>
#include<iostream>
#include<cmath>
#include<algorithm>
#include<cstring>
#include<cstdlib>
#include<cctype>
#include<vector>
#include<stack>
#include<queue>
using namespace std;
#define enter puts("") 
#define space putchar(' ')
#define Mem(a, x) memset(a, x, sizeof(a))
#define In inline
typedef long long ll;
typedef double db;
const int INF = 0x3f3f3f3f;
const db eps = 1e-8;
const int maxn = 1e5 + 5;
const int N = 18;
inline ll read()
{
  ll ans = 0;
  char ch = getchar(), last = ' ';
  while(!isdigit(ch)) last = ch, ch = getchar();
  while(isdigit(ch)) ans = (ans << 1) + (ans << 3) + ch - '0', ch = getchar();
  if(last == '-') ans = -ans;
  return ans;
}
inline void write(ll x)
{
  if(x < 0) x = -x, putchar('-');
  if(x >= 10) write(x / 10);
  putchar(x % 10 + '0');
}

char s[10];
int n;
struct Edge
{
  int nxt, to, w;
}e[maxn << 1];
int head[maxn], ecnt = -1;
In void addEdge(int x, int y, int w)
{
  e[++ecnt] = (Edge){head[x], y, w};
  head[x] = ecnt;
}

int dep[maxn], fa[N + 2][maxn];
ll dis[maxn];
In void dfs(int now, int _f)
{
  for(int i = 1; i <= N; ++i) fa[i][now] = 0;
  for(int i = 1; (1 << i) <= dep[now]; ++i)
    fa[i][now] = fa[i - 1][fa[i - 1][now]];
  for(int i = head[now], v; ~i; i = e[i].nxt)
    {
      if((v = e[i].to) == _f) continue;
      dep[v] = dep[now] + 1;
      fa[0][v] = now;
      dis[v] = dis[now] + e[i].w;
      dfs(v, now);
    }
}

In int lca(int x, int y)
{
  if(dep[x] < dep[y]) swap(x, y);
  for(int i = N; i >= 0; --i)
    if(dep[x] - (1 << i) >= dep[y]) x = fa[i][x];
  if(x == y) return x;
  for(int i = N; i >= 0; --i)
    if(fa[i][x] ^ fa[i][y]) x = fa[i][x], y = fa[i][y];
  return fa[0][x];
}

In ll query(int x, int y)
{
  int z = lca(x, y);
  return dis[x] + dis[y] - (dis[z] << 1);
}
In ll solve(int x, int y, int k)
{
  int z = lca(x, y);
  if(dep[x] - dep[z] < k) k = dep[y] + dep[x] - (dep[z] << 1) - k, swap(x, y);
  for(int i = N; i >= 0; --i)
    if((1 << i) <= k) x = fa[i][x], k -= (1 << i);
  return x;
}

int main()
{
  int T = read();
  while(T--)
    {
      Mem(head, -1), ecnt = -1;
      n = read();
      for(int i = 1; i < n; ++i)
	{
	  int x = read(), y = read(), w = read();
	  addEdge(x, y, w), addEdge(y, x, w);
	}
      dfs(1, 0);
      scanf("%s", s);
      while(s[1] != 'O')
	{
	  int x = read(), y = read();
	  if(s[1] == 'I') write(query(x, y)), enter;
	  else
	    {
	      int k = read();
	      write(solve(x, y, k - 1)), enter;
	    }
	  scanf("%s", s);
	}
    } 
  return 0;
}