Luogu P4178 Tree （点分治 + 树状数组）

Posted 2021-08-27 繁凡さん

整理的算法模板合集： ACM模板

点我看算法全家桶系列！！！

实际上是一个全新的精炼模板整合计划

---

Weblink

https://www.luogu.com.cn/problem/P4178

Problem

给定一棵 $n$ 个节点的树，每条边有边权，求出树上两点距离小于等于 $k$ 的点对数量。

$1\le n\le 4\times 10^4$，$1\le u,v\le n$，$0\le w\le 10^3$ ，$0\le k\le 2\times 10^4$。

Solution

树上路径信息问题显然用点分治。

这道题统计的是树上两点距离小于等于 $k$ 的点对数量，我们当然可以先算出距离等于 $k$ 的点对数量，然后容斥即可。

但我不想动脑子（

我们显然可以开一个桶 bit[i] 表示距离小于等于 i 的点对的数量。

我们使用树状数组维护即可。

时间复杂度 $O(n\log n\log k)$

具体的，我们先找到分治点 $root$，然后计算与 $root$ 直连的结点 $u$ 为根的子树距 $root$ 的距离， 与已经计算好了的其他子树距离去匹配，这棵子树一半，另一颗子树一半，拼起来 dis <= k 即可。因此我们需要先计算对 ans 的贡献，再更新树状数组。 最后因为我们统计完以 $u$ 为根的子树的贡献之后，会再去统计与 $root$ 直连的其他结点的贡献，所以我们需要把这棵子树的更新删除，我们记录该子树所有的点，$O(n)$ 删除即可。

最后画图可知，所有直连的结点 $u$ 到根节点 $root$ 的距离为 dis[u] = edge[root -> u]，若 dis[u] <= k ，显然点对 $(u,root)$ 也是一组答案， 并且这个时候才有资格对 ans 有 query(k - dis[u]) 的贡献，所以最后的答案 ans += query(k - dis[v]) + 1 ，要加上点对 $(u,root)$ 的贡献 $1$。

Code

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int N = 5e5 + 6, M = 3e6 + 7, INF = 0x3f3f3f3f;

int n, m, s, t, k;
int head[N], ver[M], nex[M], edge[M], tot;
bool vis[N];
ll ans;
int dis[N];
int sta[N], num1;
int seq[N], num2;
int root, siz[N], min_son_size; 
ll bit[N << 1]; 
#define lowbit(x) (x & (-x)) 

void update(int x, int val) 
{
	for (int i = x; i <= N - 7; i += lowbit(i))
		bit[i] += val;
}

ll query(int x)
{
	ll ans = 0;
	for (int i = x; i; i -= lowbit(i)) 
		ans += bit[i];
	return ans;
}

void add(int x, int y, int z)
{
	ver[tot] = y;
	edge[tot] = z;
	nex[tot] = head[x];
	head[x] = tot ++ ;
}

void get_rt(int x, int fa, int n)
{
	siz[x] = 1;
	int tmp = 0;
	for (int i = head[x]; ~i; i = nex[i]) {
		int y = ver[i];
		if(y == fa || vis[y]) continue;
		siz[x] += siz[y];
		tmp = max(tmp, siz[y]);
	}
	tmp = max(tmp, n - siz[x]);
	if(tmp < min_son_size)
		min_son_size = tmp, root = x;
}

void get_dis(int x, int fa)
{
	sta[ ++ num1] = x;
	for (int i = head[x]; ~i; i = nex[i]) {
		int y = ver[i], z = edge[i];
		if(y == fa || vis[y]) continue;
		dis[y] = dis[x] + z;
		get_dis(y, x);
	}
} 

void cal(int x)
{
	num2 = 0; 
	for (int i = head[x]; ~i; i = nex[i]) {
		int y = ver[i], z = edge[i];
		if(vis[y]) continue;
		num1 = 0;
		dis[y] = z;
		get_dis(y, x);
		for (int j = 1; j <= num1; ++ j) {
			if(dis[sta[j]] <= k) {
				ans ++ ;
				ans += query(k - dis[sta[j]]);
			}
		}
		for (int j = 1; j <= num1; ++ j) {
			if(dis[sta[j]] <= k) {
				update(dis[sta[j]], 1);
				seq[ ++ num2] = sta[j];
			}
		}
	}
	for (int i = 1; i <= num2; ++ i) {
		int pos = seq[i];
		update(dis[seq[i]], -1);
		dis[pos] = 0;
	}
}

void divide(int x)
{
	vis[x] = 1;
	cal(x);
	for (int i = head[x]; ~i; i = nex[i]) {
		int y = ver[i];
		if(vis[y]) continue;
		min_son_size = INF;
		get_rt(y, x, siz[y]);
		divide(root);
	}
}

int main()
{
	memset(head, -1, sizeof head);
	scanf("%d", &n);
	for (int i = 1; i <= n - 1; ++ i) {
		int u, v, w;
		scanf("%d%d%d", &u, &v, &w);
		add(u, v, w);
		add(v, u, w);
	}
	scanf("%d", &k);

	min_son_size = INF;
	get_rt(1, 0, n);
	divide(root);
	
	cout << ans << endl;
	return 0;
}