SP913 QTREE2 - Query on a tree II

原题传送门

题目大意

给定一棵 $n$ 个点的树，边具有边权。要求作以下操作：

• DIST：询问点 $a$ 至点 $b$ 路径上的边权之和。

• KTH：询问点 $a$ 至点 $b$ 有向路径上的第 $k$ 个点的编号。

有多组测试数据，每组数据以 DONE 结尾。

解题思路

前置知识：树上差分，树上倍增，最近公共祖先。

设 $lca$ 为 $u$ 和 $v$ 的最近公共祖先。

对于每个 DIST 操作，答案显然为 $dis[u] + dis[v] - 2 * dis[lca]$（典型的树上边差分）。

对于每个 KTH 操作，分情况讨论。

• 若 $dep[u] - dep[lca] + 1 < k$，答案即为 $u$ 的 $k$ 级祖先；
• 若 $dep[u] - dep[lca] + 1 > k$，答案为 $v$ 的 $dep[u] + dep[v] - 2 * dep[k] + 2 - k$ 级祖先；
• 若 $dep[u] - dep[lca] + 1 = k$，答案为 $lca$。

时间复杂度：$O((n + q) \log n)$。

AC CODE

笔者快读太长，不便放在这里。

#include<bits/stdc++.h>
using namespace std;
#define _ (int) 1e4 + 5

int T, n;

array<int, _> val;

int tot;
array<int, _> head;
array<int, _ << 1 > to, nxt, w;

array<int, _> dep;
array<array<int, 21>, _> fa;

void add(int a, int b, int c)
{
	to[++tot] = b;
	w[tot] = c;
	nxt[tot] = head[a];
	head[a] = tot;
}

void dfs(int u, int f, int d = 1)
{
	fa[u][0] = f;
	dep[u] = d;
	for(int  i = head[u]; i; i = nxt[i])
	{
		int v = to[i];
		if(v == f) continue;
		val[v] = val[u] + w[i];
		dfs(v, u, d + 1);
	}
}

int LCA(int u, int v)
{
	if(dep[u] > dep[v]) swap(u, v);
	for(int i = 20; i >= 0; --i)
		if(dep[fa[v][i]] >= dep[u])
			v = fa[v][i];
	if(u == v) return u;
	for(int i = 20; i >= 0; --i)
		if(fa[u][i] != fa[v][i])
		{
			u = fa[u][i];
			v = fa[v][i];
		}
	return fa[u][0];
}

int query(int u, int v, int k)
{
	int f = LCA(u, v);
	if (dep[u] - dep[f] + 1 <= k)
	{
		k = dep[u] + dep[v] - (dep[f] << 1) + 2 - k;
		u = v;
	}
	for (int i = 20; i >= 0; --i)
		if (k >= (1 << i) + 1)
		{
			u = fa[u][i];
			k -= (1 << i);
		}
	return u;
}

void init()
{
	tot = 0;
	head.fill(0);
	dep.fill(0);
	fa.fill({0, 0});
	val.fill(0);
}

signed main()
{
	cin >> T;
	while(T--)
	{
		init();
  		cin >> n;
		for(int i = 1; i < n; ++i)
		{
			int a, b, c;
			cin >> a >> b >> c;
			add(a, b, c);
			add(b, a, c);
		}

		dfs(1, 0);
		for(int j = 1; j <= 20; ++j)
			for(int i = 1; i <= n; ++i)
				fa[i][j] = fa[fa[i][j - 1]][j - 1];

		while(1)
		{
			char c[10];
			int x, y, z;
			scanf("%s", c);
			if(c[1] == 'O') break;
			if(c[1] == 'I')
			{
				cin >> x >> y;
				cout << val[x] + val[y] - 2 * val[LCA(x, y)] << endl;
			}
			if(c[1] == 'T')
			{
				cin >> x >> y >> z;
				cout << query(x, y, z) << endl;
			}
		}
	}
}

/*
1
6
1 2 1
2 4 1
2 5 2
1 3 1
3 6 2
DIST 4 6
KTH 4 6 4
DONE
*/