男人八题-点分治-Acwing252.树

Acwing252.树

题目分析

树中的路径分为三种

• 路径两端在同一个子树
• 路径两端在不同子树
• 路径有一端是重心

因此可以分情况处理, 对于第一种情况可以进行递归处理, 第二种情况需要使用容斥原理求得(下面重点介绍), 第三种情况枚举重心到其他节点的路径就可以求得

代码分析

求子树大小

int get_size(int _index, int pre) {
//	如果当前子树被删除, 返回0
	if (visited[_index]) return 0;

	int res = 1;
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		if (ver == pre) continue;
		res += get_size(ver, _index);
	}

	return res;
}

求当前点出发的子树的所有路径

void get_dist(int _index, int pre, int dis, int &ptr) {
	if (visited[_index]) return;
	arr2[ptr++] = dis;
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		if (ver == pre) continue;
		get_dist(ver, _index, dis + w[i], ptr);
	}
}

求树的重心

int get_weight_center(int _index, int pre, int tot, int &res) {
	if (visited[_index]) return 0;
	int sum = 1, max_son = 0;

	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		if (ver == pre) continue;
		int size = get_weight_center(ver, _index, tot, res);
//		更新最大子树大小
		max_son = max(max_son, size);
		sum += size;
	}

//	总数减下面的子树总和, 就是当前点所在父节点的子树大小
	max_son = max(max_son, tot - sum);
//	如果最大子树大小 <= tot / 2, 那么当前点可以作为重心
	if (max_son <= tot >> 1) res = _index;
	return sum;
}

注意:这里树的重心只要子树的大小<= n / 2就可以保证点分治时间复杂度

求单个子树中路径和 <= k的路径组合数量

int get(int arr[], int cnt) {
	sort(arr, arr + cnt);
	int res = 0;

	for (int i = cnt - 1, j = -1; i >= 0; --i) {
//		找到满足条件的最大j
		while (j + 1 < i && arr[j + 1] + arr[i] <= k) j++;
		j = min(j, i - 1);
		res += j + 1;
	}

	return res;
}

双指针算法求解

核心函数

int calc(int _index) {
	if (visited[_index]) return 0;
	int res = 0;

//	找到当前子树的重心
	get_weight_center(_index, -1, get_size(_index, -1), _index);
	visited[_index] = true;

	int p1 = 0;
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i], p2 = 0;
		get_dist(ver, -1, w[i], p2);
//		减去一个子树中的组合情况
		res -= get(arr2, p2);

//		将子树的距离加入到去全局的距离中
		for (int j = 0; j < p2; ++j) {
//			某个点是重心的情况
			if (arr2[j] <= k) res++;
			arr1[p1++] = arr2[j];
		}
	}

//	加上子树内部的组合情况
	res += get(arr1, p1);
//	递归处理子树
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		res += calc(ver);
	}
	return res;
}

先找到当前子树的重心, 然后将重心标记

然后将路径分为三种情况, 上面有论述,
直接看最复杂的情况, 也就是两个点分别在两个子树中

• 求得子树中的路径组合, 因为是不合法的, 因此将其从答案中减去
• 然后将子树的路径添加到全局的路径当中
• 在添加过程中不要忘记记录子树到重心也是一种情况
• 最后加入到res的就是全局的除去子树之间的路径组合

最后处理第一种和第三种情况

时间复杂度分析

递归过程中一共$\log{}{n}$层, 每一层点数不超过$n$, 排序时间复杂度$O(n\log_{}{n})$, 因此总的时间复杂度$O(n\log_{}{n}^{2})$

完整代码

#include <iostream>
#include <algorithm>
#include <cstring>

using namespace std;

const int N = 10010, EDGE_NUMBER = N << 1;

int n, k;
int head[N], edge_end[EDGE_NUMBER], next_edge[EDGE_NUMBER], w[EDGE_NUMBER], edge_index;
//节点是否被删除
bool visited[N];
//当前重心的所有子树的距离, 当前子树的距离
int arr1[N], arr2[N];

void add(int ver1, int ver2, int val) {
	edge_end[edge_index] = ver2, next_edge[edge_index] = head[ver1], w[edge_index] = val, head[ver1] = edge_index++;
}

//求子树大小
int get_size(int _index, int pre) {
//	如果当前子树被删除, 返回0
	if (visited[_index]) return 0;

	int res = 1;
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		if (ver == pre) continue;
		res += get_size(ver, _index);
	}

	return res;
}

//求从当前点出发所有路径的长度
void get_dist(int _index, int pre, int dis, int &ptr) {
	if (visited[_index]) return;
	arr2[ptr++] = dis;
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		if (ver == pre) continue;
		get_dist(ver, _index, dis + w[i], ptr);
	}
}

//求树的重心, 返回值是当前子树大小
int get_weight_center(int _index, int pre, int tot, int &res) {
	if (visited[_index]) return 0;
	int sum = 1, max_son = 0;

	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		if (ver == pre) continue;
		int size = get_weight_center(ver, _index, tot, res);
//		更新最大子树大小
		max_son = max(max_son, size);
		sum += size;
	}

//	总数减下面的子树总和, 就是当前点所在父节点的子树大小
	max_son = max(max_son, tot - sum);
//	如果最大子树大小 <= tot / 2, 那么当前点可以作为重心
	if (max_son <= tot >> 1) res = _index;
	return sum;
}

int get(int arr[], int cnt) {
	sort(arr, arr + cnt);
	int res = 0;

	for (int i = cnt - 1, j = -1; i >= 0; --i) {
//		找到满足条件的最大j
		while (j + 1 < i && arr[j + 1] + arr[i] <= k) j++;
		j = min(j, i - 1);
		res += j + 1;
	}

	return res;
}

int calc(int _index) {
	if (visited[_index]) return 0;
	int res = 0;

//	找到当前子树的重心
	get_weight_center(_index, -1, get_size(_index, -1), _index);
	visited[_index] = true;

	int p1 = 0;
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i], p2 = 0;
		get_dist(ver, -1, w[i], p2);
//		减去一个子树中的组合情况
		res -= get(arr2, p2);

//		将子树的距离加入到去全局的距离中
		for (int j = 0; j < p2; ++j) {
//			某个点是重心的情况
			if (arr2[j] <= k) res++;
			arr1[p1++] = arr2[j];
		}
	}

//	加上子树内部的组合情况
	res += get(arr1, p1);
//	递归处理子树
	for (int i = head[_index]; ~i; i = next_edge[i]) {
		int ver = edge_end[i];
		res += calc(ver);
	}
	return res;
}

int main() {
	while (scanf("%d%d", &n, &k), n && k) {
		memset(visited, false, sizeof visited);
		memset(head, -1, sizeof head);
		edge_index = 0;

		for (int i = 0; i < n - 1; ++i) {
			int ver1, ver2, val;
			scanf("%d%d%d", &ver1, &ver2, &val);
			add(ver1, ver2, val);
			add(ver2, ver1, val);
		}

		printf("%d\n", calc(0));
	}
	return 0;
}