树上询问

$$\mathrm{树上询问}$$

题目链接：$\mathrm{l}\mathrm{u}\mathrm{o}\mathrm{g}\mathrm{u}\mathrm{P}6374$

题目

给定一棵 $n$ 个点的无根树，有 $q$ 次询问。

每次询问给一个参数三元组 $(a,b,c)$ ，求有多少个 $i$ 满足这棵树在以 $i$ 为根的情况下 $a$ 和 $b$ 的 LCA 为 $c$ 。

输入

第一行 $2$ 个数，为 $n$ 和 $q$ 。

接下来 $n-1$ 行，每行 $2$ 个数，表示树的一条边。

接下来 $q$ 行，每行 $3$ 个数，为 $(a,b,c)$ 。

输出

共 $q$ 行，每行一个数，为对于每个三元组的 $i$ 的个数。

样例输入1

10 5
1 2
1 3
2 4
2 5
2 10
5 6
3 7
7 8
7 9
4 6 2
4 10 1
6 8 3
9 10 2
4 10 5

样例输出1

7
0
1
4
0

样例输入2

5 3
1 3
1 5
3 4
3 2
5 2 3
5 2 1
2 4 5

样例输出2

2
1
0

样例解释2

第一个查询的 $i$ 为 $3$ 和 $4$ 。

第二个查询的 $i$ 为 $1$ 。

样例输入3

20 10
1 2
1 3
1 4
2 5
2 6
3 10
4 13
4 14
6 7
6 8
10 11
4 15
4 16
8 9
11 12
16 17
16 18
16 19
17 20
15 19 16
1 12 1
20 20 20
7 7 8
1 8 3
5 20 2
2 9 6
9 12 1
9 12 2
9 12 3

样例输出3

4
16
20
0
0
5
2
10
2
1

数据范围

$subtask1(20pts)$ ： $1\le n\le 1000$ ， $1\le q\le 500$ 。

$subtask2(15pts)$ ： $1\le n\le 10^5$ ， $1\le q\le 10^5$ ，树退化成链 。

$subtask3(25pts)$ ： $1\le n\le 5\times 10^5$ ， $1\le q\le 10^5$ ，数据不随机 。

$subtask4(40pts)$ ： $1\le n\le 5\times 10^5$ ， $1\le q\le 2\times 10^5$ 。

对于所有数据： $1\le n\le 5\times 10^5$ ， $1\le q\le 2\times 10^5$ 。

注：数据强度不高，不必卡常与快读快输。

思路

这道题是一道 LCA 。

我们先让根节点是 $1$ ，然后对于每一个三元组，就先求出 $1$ 是根节点的时候， $x$ 和 $y$ 的 LCA （设为 $lca$ ）。那我们通过找规律（这里不说发现过程，直接看这位大佬的博客），可以发现就是当 $z$ 为根节点时，除了 $x,y$ 所在的子树的节点，其它节点都可成为根节点。
那就有四种情况：

1. 刚好就是 $lca=z$ ，答案就是 $n-(x,y$ 所在的子树的节点 $)$
2. $z$ 在 $x$ 到 $lca$ 的路径上，答案是 $n-(x$ 和不在 $z$ 所在的子树的节点 $)$
3. $y$ 在 $x$ 到 $lca$ 的路径上，答案是 $n-(y$ 和不在 $z$ 所在的子树的节点 $)$
4. $z$ 不在 $x$ 到 $y$ 的路径上，就没有方案，答案是 $0$ 。

就看是哪个，输出就可以了。

代码

#include<cmath>
#include<cstdio>
#include<algorithm>

using namespace std;

struct node {
	int to, nxt;
}e[1000001];
int n, q, x, y, z, KK, le[500001], fa[21][500001], num[500001], dep[500001], lca;

void add(int x, int y) {
	e[++KK] = (node){y, le[x]}; le[x] = KK;
	e[++KK] = (node){x, le[y]}; le[y] = KK;
}

void dfs(int now, int father) {//遍历初始化
	fa[0][now] = father;
	dep[now] = dep[father] + 1;
	for (int i = le[now]; i; i = e[i].nxt)
		if (e[i].to != father) {
			dfs(e[i].to, now);
			num[now] += num[e[i].to];
		}
	num[now]++;
}

int LCA(int x, int y) {//LCA
	if (dep[x] < dep[y]) swap(x, y);
	for (int i = log2(n) + 1; i >= 0; i--)
		if (dep[fa[i][x]] >= dep[y])
			x = fa[i][x];
	if (x == y) return x;
	for (int i = log2(n) + 1; i >= 0; i--)
		if (fa[i][x] != fa[i][y]) {
			x = fa[i][x];
			y = fa[i][y];
		}
	return fa[0][x];
}

int get_sum(int x, int y) {//得到y点包含x的子树的大小
	if (x == y) return 0;
	for (int i = log2(n) + 1; i >= 0; i--)
		if (dep[fa[i][x]] > dep[y])
			x = fa[i][x];
	return num[x];
}

bool ch(int x, int y, int z) {//判断z是否在两个点的lca和其中一个点中间
	int k = LCA(x, z);
	return ((LCA(x, y) == y) || (LCA(y, z) == y)) && (LCA(y, k) == k);
}

int main() {
	scanf("%d %d", &n, &q);//读入
	for (int i = 1; i < n; i++) {
		scanf("%d %d", &x, &y);
		add(x, y);//建图
	}
	
	dfs(1, 0);//遍历初始化
	
	for (int i = 1; i <= 20; i++)//初始化父亲
		for (int j = 1; j <= n; j++)
			fa[i][j] = fa[i - 1][fa[i - 1][j]];
	
	for (int i = 1; i <= q; i++) {
		scanf("%d %d %d", &x, &y, &z);//读入
		lca = LCA(x, y);//求出LCA
		if (lca == z) {//LCA就是
			printf("%d\n", n - get_sum(x, z) - get_sum(y, z));
			continue;
		}
		if (ch(x, z, lca)) {//在x和LCA之间
			printf("%d\n", n - get_sum(x, z) - (n - num[z]));
			continue;
		}
		if (ch(y, z, lca)) {//在y和LCA之间
			printf("%d\n", n - get_sum(y, z) - (n - num[z]));
			continue;
		}
		printf("0\n");//不在x和y的路上，莫得方案
	}
	
	return 0;
}