QTREE2 - Query on a tree II

简单倍增。

前置：树上倍增 k 级祖先。

注意题目无修改，所以对于询问距离，直接预处理出每个点到根的距离，然后相加减去两倍 LCA 即可，非常套路。

对于第二位，先求出 $u,v$ 的 LCA，先看 $u \rightarrow lca$ 这条路径上经过的点数是否 $\geq k$，若 $\geq k$，求 $u$ 的 $k$ 级祖先即可。若点数 $<k$，从 $lca$ 开始往下遍历，可以转化为 $v$ 的多级祖先，倍增处理即可。

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

const int N = 1e4 + 5;

int fa[N][31], sum1[N], sum2[N];
vector<pair<int, int> > G[N];
int n, t;

class TreeCut
{
public:
	int fa[N], sz[N], son[N], top[N], dep[N];
	void Init(int n)
	{
		for (int i = 0; i <= n; i++)
		{
			fa[i] = sz[i] = son[i] = top[i] = dep[i] = 0;
		}
	}
	void dfs1(int u, int f)
	{
		fa[u] = f;
		dep[u] = dep[f] + 1;
		sz[u] = 1;
		for (auto j : G[u])
		{
			if (j.first != f)
			{
				dfs1(j.first, u);
				sz[u] += sz[j.first];
				if (sz[son[u]] < sz[j.first]) son[u] = j.first;
			}
		}
	}
	void dfs2(int u, int tf)
	{
		top[u] = tf;
		if (!son[u]) return;
		dfs2(son[u], tf);
		for (auto j : G[u])
		{
			if (j.first != son[u] && j.first != fa[u])
			{
				dfs2(j.first, j.first);
			}
		}
	}
	int lca(int u, int v)
	{
		while (top[u] ^ top[v])
		{
			if (dep[top[u]] < dep[top[v]])
			{
				u ^= v;
				v ^= u;
				u ^= v;
			}
			u = fa[top[u]];
		}
		if (dep[u] < dep[v]) swap(u, v);
		return v;
	}
}f;

void dfs(int u, int f, int w)
{
	fa[u][0] = f;
	sum1[u] = sum1[f] + w;
	sum2[u] = sum2[f] + 1;
	for (auto j : G[u])
	{
		if (j.first != f) dfs(j.first, u, j.second);
	}
}

void solve_fa()
{
	for (int i = 1; i <= 30; i++)
	{
		for (int j = 1; j <= n; j++)
		{
			fa[j][i] = fa[fa[j][i - 1]][i - 1];
		}
	}
}

int main()
{
	ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);
	cin >> t;
	while (t--)
	{
		cin >> n;
		for (int i = 1; i <= n; i++) G[i].clear();
		for (int i = 1; i < n; i++)
		{
			int u, v, w;
			cin >> u >> v >> w;
			G[u].push_back(make_pair(v, w));
			G[v].push_back(make_pair(u, w));
		}
		for (int i = 1; i <= n; i++)
		{
			for (int j = 0; j <= 30; j++) fa[i][j] = 0;
		}
		f.Init(n);
		f.dfs1(1, 0);
		f.dfs2(1, 1);
		dfs(1, 0, 0);
		solve_fa();
		string s;
		while (cin >> s && s != "DONE")
		{
			if (s[0] == 'D')
			{
				int u, v;
				cin >> u >> v;
				int lca = f.lca(u, v);
				cout << sum1[u] + sum1[v] - 2 * sum1[lca] << "\n";
			}
			else
			{
				int u, v, k;
				cin >> u >> v >> k;
				int lca = f.lca(u, v);
				int dist_of_u_and_lca = sum2[u] - sum2[lca] + 1;
				if (dist_of_u_and_lca >= k)
				{
					// 所求的点在 u 到 lca 的路径上
					int now = u, p = 0;
					k--;
					while (k)
					{
						if (k & 1)
						{
							now = fa[now][p];
						}
						k >>= 1;
						p++;
					}
					cout << now << "\n";
				}
				else
				{
					//cout << dist_of_u_and_lca << "\n";
					k -= dist_of_u_and_lca;
					int dist_of_v_and_lca = sum2[v] - sum2[lca];
					//cout << dist_of_v_and_lca << "\n";
					k = dist_of_v_and_lca - k;
					int now = v, p = 0;
					//k--;
					while (k)
					{
						if (k & 1)
						{
							now = fa[now][p];
						}
						p++;
						k >>= 1;
					}
					cout << now << "\n";
				}
			}
		}
	}
	return 0;
}