P5838 [USACO19DEC]Milk Visits G (树链剖分 + 主席树)

题目链接: P5838 [USACO19DEC]Milk Visits G

大致题意

给定一棵有$n$个节点的树, 每个节点有权值$w_i$, 有$m$次询问:

a b c 询问$a,b$路径上是否有$w_i==c$的节点.

解题思路

树链剖分 + 主席树

我们首先考虑如果不是树上路径问题, 而是询问$[l,r]$区间中, 是否有$w_i==c$的点.

我们可以采用可持久化线段树来实现.

对于树上路径问题, 我们加上树链剖分即可.

---

特别的, 建树应当对于$dfs$序去建树, 每次树上路径查询时, 同样需要对左端点去$-1$.

AC代码

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 1; i <= (n); ++i)
using namespace std;
typedef long long ll;
const int N = 1E5 + 10;
int n, m;
int w[N]; //给定的树中各个顶点的权值
vector<int> edge[N]; //树上各个点之间的边

int p[N], dep[N], sz[N], son[N];
// 父节点   深度   节点大小 重儿子
void dfs1(int x = 1, int fa = 0) { // x = 树根节点
	p[x] = fa, dep[x] = dep[fa] + 1, sz[x] = 1; // son[x] = 0;
	for (auto& to : edge[x]) {
		if (to == fa) continue;
		dfs1(to, x);
		sz[x] += sz[to];		// 特别的, 如果边权->点权, 应记录w[to] = 边权.
		if (sz[to] > sz[son[x]]) son[x] = to; //更新重儿子
	}
}
int id[N], nw[N], top[N], ind;
//  新编号  新值    重链顶 当前用到的编号
void dfs2(int x = 1, int tp = 1) { // x = 树根节点, tp = 树根节点
	id[x] = ++ind, nw[ind] = w[x], top[x] = tp;

	if (!son[x]) return; //叶子结点
	dfs2(son[x], tp); //先遍历重儿子

	for (auto& to : edge[x]) {
		if (to == p[x] or to == son[x]) continue;
		dfs2(to, to);
	}
}


namespace {
	struct node {
		int l, r;
		int cou;
	}t[N * 18];
	int root[N], ind;
	int build(int a, int c, int tl, int tr, int p) {
		int x = ++ind; t[x] = t[p];
		t[x].cou += c;
		if (tl == tr) return x;
		int mid = tl + tr >> 1;
		if (a <= mid) t[x].l = build(a, c, tl, mid, t[p].l);
		else t[x].r = build(a, c, mid + 1, tr, t[p].r);
		return x;
	}
	int ask(int a, int tl, int tr, int p, int x) {
		if (tl == tr) return t[x].cou - t[p].cou;
		int mid = tl + tr >> 1;
		if (a <= mid) return ask(a, tl, mid, t[p].l, t[x].l);
		return ask(a, mid + 1, tr, t[p].r, t[x].r);
	}
}

bool ask_route(int a, int b, int c) {
	while (top[a] != top[b]) {
		if (dep[top[a]] < dep[top[b]]) swap(a, b);
		int r = id[a], l = id[top[a]];
		if (ask(c, 1, n, root[l - 1], root[r])) return 1;
		a = p[top[a]];
	}

	if (dep[a] > dep[b]) swap(a, b);
	return ask(c, 1, n, root[id[a] - 1], root[id[b]]);
}
char res[N];
int main()
{
	cin >> n >> m;
	rep(i, n) scanf("%d", &w[i]);

	rep(i, n - 1) {
		int a, b; scanf("%d %d", &a, &b);
		edge[a].push_back(b), edge[b].push_back(a);
	}

	dfs1(), dfs2();
	rep(i, n) root[i] = build(nw[i], 1, 1, n, root[i - 1]);

	rep(i, m) {
		int a, b, c; scanf("%d %d %d", &a, &b, &c);
		res[i] = '0' + ask_route(a, b, c);
	}
	res[m + 1] = 0;
	puts(res + 1);

	return 0;
}

END