2019 ICPC 南昌 K-tree【树上启发式合并+线段树】

最新推荐文章于 2022-03-28 18:57:48 发布

原创最新推荐文章于 2022-03-28 18:57:48 发布 · 307 阅读

1 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#数据结构

数据结构同时被 2 个专栏收录

46 篇文章

订阅专栏

题解

36 篇文章

订阅专栏

这篇博客介绍了如何利用Disjoint Set Union (DSU) 和线段树来解决一类树上的问题。给定一棵有根树，每个节点带有权值，目标是找出所有满足特定条件的节点对(u, v)，即u和v不在同一子树内且u的权值加v的权值等于两倍它们最近公共祖先的权值，并且u和v的深度之和减去最近公共祖先的深度的两倍小于等于K。博主提供了详细的算法实现，包括DSU的启发式合并和线段树的更新与查询操作，最后给出了完整的C++代码实现。

传送门

分析

dsu on tree + 线段树

给定一个有根树，每个节点有一个权值
找到有多少不同的对 $(u, v)$
满足条件

$u, v$ 任何一个都不是另一个的祖先节点，也就是说， $u, v$ 在不同的子树里
$u_{value}+v_{value}=2*lca(u, v)_{value}$
$dep_u + dep_v - 2*dep_{lca} <= K$

一眼树上启发式合并？
dsu直接套上去，用值域线段树维护子节点信息，算答案的时候
枚举 $l c a$ 的同时，枚举v，其中v为较小的那个子树集合里面的点，保证复杂度
知道 $l c a$ 和 $v$ ，就能通过公式确定上面的值范围
在权值线段树中找那个值范围节点的个数即可

代码如下

代码

//t42586 
/*
  @Author: YooQ
*/
#include <bits/stdc++.h>
using namespace std;
#define sc scanf
#define pr printf
#define ll long long
#define FILE_OUT freopen("out", "w", stdout);
#define FILE_IN freopen("in", "r", stdin);
#define debug(x) cout << #x << ": " << x << "\n";
#define AC 0
#define WA 1
#define INF 0x3f3f3f3f
const ll MAX_N = 2e5+5;
const ll MOD = 1e9+7;
int N, M, K;

int arr[MAX_N];
int head[MAX_N];
int tot = 0;
struct Edge{
	int to, nxt;
}edge[MAX_N];

void addEdge(int u, int v) {
	edge[tot].nxt = head[u];
	edge[tot].to = v;
	head[u] = tot++;
}

struct Tr {
	int k, l, r;
}tr[MAX_N<<5];
int root[MAX_N];
int indx = 0;

int mk() {
	return ++indx;
}

void push_up(int rt) {
	tr[rt].k = tr[tr[rt].l].k + tr[tr[rt].r].k;
}

void update(int& rt, int l, int r, int x, int k) {
	if (!rt) rt = mk();
	if (l == r) {
		tr[rt].k += k;
		return;
	}
	int mid = l + ((r-l)>>1);
	if (x <= mid) update(tr[rt].l, l, mid, x, k);
	if (x  > mid) update(tr[rt].r, mid+1, r, x, k);
	push_up(rt);
}

int query(int rt, int l, int r, int x, int y) {
	if (!rt || l > r) return 0;
	if (x <= l && r <= y) {
		return tr[rt].k;
	}
	int mid = l + ((r-l)>>1);
	if (y <= mid) return query(tr[rt].l, l, mid, x, y);
	if (x  > mid) return query(tr[rt].r, mid+1, r, x, y);
	return query(tr[rt].l, l, mid, x, y) + query(tr[rt].r, mid+1, r, x, y);
}

int sz[MAX_N];
int son[MAX_N];
int dep[MAX_N];

void dfs(int u, int d) {
	dep[u] = d;
	sz[u] = 1;
	son[u] = 0;
	int v;
	for (int i = head[u];~i;i=edge[i].nxt) {
		dfs(v=edge[i].to, d+1);
		sz[u] += sz[v];
		if (sz[v] > sz[son[u]]) {
			son[u] = v;
		}
	}
}

void upd(int u, int k) {
	int d = dep[u];
	update(root[arr[u]], 1, N, d, k);
}

void calc(int u) {
	upd(u, 1);
	for (int i = head[u];~i;i=edge[i].nxt) {
		calc(edge[i].to);
	}
}

void del(int u) {
	upd(u, -1);
	for (int i = head[u];~i;i=edge[i].nxt) {
		del(edge[i].to);
	}
}
ll ans = 0;

void get_ans(int u, int lca) {
	int limit = K + dep[lca] * 2 - dep[u];
	int k = arr[lca] * 2 - arr[u];
	ans += query(root[k], 1, N, dep[lca]+1, min(max(limit, 0), N));
	int v;
	for (int i = head[u];~i;i=edge[i].nxt) {
		get_ans(edge[i].to, lca);
	}
}

void dsu(int u, int opt) {
	int v;
	for (int i = head[u];~i;i=edge[i].nxt) {
		if ((v=edge[i].to) == son[u]) continue;
		dsu(v, 1);
	}
	if (son[u]) dsu(son[u], 0);
	for (int i = head[u];~i;i=edge[i].nxt) {
		if ((v=edge[i].to) == son[u]) continue;
		get_ans(v, u);
		calc(v);
	}
	upd(u, 1);
	if (opt) del(u);
}


void init() {
	memset(head, -1, sizeof head);
	tot = 0;
}

void solve(){
	init();
	sc("%d%d", &N, &K);
	for (int i = 1; i <= N; ++i) {
		sc("%d", &arr[i]);
	}
	int u, v;
	for (int i = 2; i <= N; ++i) {
		sc("%d", &u);
		addEdge(u, i);
	}
	dfs(1, 1);
	dsu(1, 0);
	pr("%lld", ans*2);
}

signed main()
{
	#ifndef ONLINE_JUDGE
	//FILE_IN
	FILE_OUT
	#endif
	int T = 1;//cin >> T;
	while (T--) solve();

	return AC;
}