AcWing 252. 树——点分治

题目链接

题目大意

给你一棵N个点的树，每条边都有一个权值，树上两个点的路径长度就是路径上的权值之和，问你路径长度不超过K的路径有多少条

解题思路

这是一棵无根树，我们考虑使用点分治的方法来解决
首先，我们任选一个节点作为根节点，这样，符合答案的路径有一下两种情况
1、经过这个根节点
2、不经过这个根节点
我们可以这样考虑，我们每次只算经过根节点的路径，这样我们递归这个根节点的子树，这样我们一直往下递归，可以做到不重不漏地找出所有合法路径
假设我们选了一个根节点$root$，我们就需要得到一个数组$a$，$a[i]$表示这个根节点可以到达的子节点，我们还需要一个数组$d$，$d[x]$表示$x$到根节点p的路径长度。
如果$d[i]+d[j]<=k$ 是不是就是合法答案呢？显然不是，我们统计的路径是要经过根节点的，但是这样的$i,j$可能是在同一棵子树中的，所以我们还要记录一个$b$数组，$b[x]$表示节点$x$属于哪一个子树。
我们有了上述条件怎么快速求出合法的路径呢，我们将这棵树上的点（数组$a$）按$d$值排序，用两个指针 $l,r$ 分别从前后扫描$a$数组。此外，我们用一个 $num[s]$ 数组维护一个 $[l+1,r]$ 区间内满足$b[a[i]]=s$ 的位置 $i$ 的个数，于是我们当路径一端为 $x$ 时，另一端满足要求的 $y$ 的个数为 $$r-l-num[b[a[l]]$$
总而言之，我们就是执行以下四步
1、找到一个根节点 $p$
2、通过一遍 $dfs$ 找到 $b$ 数组和 $d$ 数组
3、计算这棵树的答案
4、删除节点 $p$，对 $p$ 的每棵子树进行这四步操作

上述办法我们的复杂度是 $O(TNlogN)$ 其中，$T$ 是递归的层数，这样就有一个问题，如果这是一条链，我们这样就崩了，所以我们每次选取根节点要选取树的重心。这样复杂度是 $O(Nlo{g}^{2}N)$

Code

#include <bits/stdc++.h>
#define ll long long
#define qc ios::sync_with_stdio(false); cin.tie(0);cout.tie(0)
#define fi first
#define se second
#define PII pair<int, int>
#define PLL pair<ll, ll>
#define pb push_back
using namespace std;
const int MAXN = 10007;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
const ll mod = 1e9 + 7;
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while (!isdigit(ch)){if (ch=='-') f=-1;ch=getchar();}
    while (isdigit(ch)){x=x*10+ch-48;ch=getchar();}
    return x*f;
}
int n, k;
int ans;
int head[MAXN];
int cnt;
int vis[MAXN];
int sz[MAXN];
int root, tot, num[MAXN];
int maxSize;
int d[MAXN], b[MAXN];
int a[MAXN];
struct Edge{
	int to, next, w;
}e[MAXN << 1];
void init(){
	for(int i = 0; i <= n; i++)
		head[i] = -1;
	cnt = 0;
}
void add(int u, int v, int w = 0){
	e[cnt].to = v;
	e[cnt].w = w;
	e[cnt].next = head[u];
	head[u] = cnt++;
}
void get_root(int S, int u, int f){
	sz[u] = 1;
	int maxPart = 0;
	for(int i = head[u]; ~i; i = e[i].next){
		int v = e[i].to;
		if(v == f || vis[v])
			continue;
		get_root(S, v, u);
		sz[u] += sz[v];
		maxPart = max(maxPart, sz[v]);
	}
	maxPart = max(maxPart, S - sz[u]);
	if(maxPart < maxSize){
		maxSize = maxPart;
		root = u;
	}
}
void dfs2(int u, int f){
	for(int i = head[u]; ~i; i = e[i].next){
		int v = e[i].to;
		if(vis[v] || v == f)
			continue;
		d[v] = d[u] + e[i].w;
		a[++tot] = v;
		b[v] = b[u];
		num[b[v]]++;
		dfs2(v, u);
	}
}
bool cmp(int i, int j){
	return d[i] < d[j];
}
void dfs(int S, int u){
	maxSize = S;
	get_root(S, u, -1);
	memset(d, 0, sizeof(d));
	memset(b, 0, sizeof(b));
	memset(num, 0, sizeof(num));
	tot = 0;
	vis[root] = 1;
	a[++tot] = root;
	d[root] = 0;
	b[root] = root;
	++num[root];
	for(int i = head[root]; ~i; i = e[i].next){
		int v = e[i].to;
		if(vis[v])
			continue;
		d[v] = e[i].w;
		a[++tot] = v;
		b[v] = v;
		num[v]++;
		dfs2(v, root);
	}
	sort(a+1, a+1+tot, cmp);
	int l = 1, r = tot;
	num[b[a[1]]]--;
	while(l < r){
		while(d[a[l]] + d[a[r]] > k) --num[b[a[r--]]];
		ans += r - l - num[b[a[l]]];
		l++;
		num[b[a[l]]]--;
		
	} 
	for(int i = head[root]; ~i; i = e[i].next){
		int v = e[i].to;
		if(vis[v])
			continue;
		dfs(sz[v], v);
	}
}
void solve(){
    init();
    for(int i = 1; i < n; i++){
    	int u, v, w;
    	cin >> u >> v >> w;
    	add(u, v, w);
    	add(v, u, w);
    }
    memset(sz, 0, sizeof(sz));
    memset(vis, 0, sizeof(vis));
    ans = 0;
    dfs(n, 1);
    cout << ans << endl;
}

int main()
{
	#ifdef ONLINE_JUDGE
    #else
       freopen("in.txt", "r", stdin);
       freopen("out.txt", "w", stdout);
    #endif

	qc;
    int T;
    // cin >> T;
    T = 1;
    while(T--){
    	while(cin >> n >> k && (n+k))
        solve();
    }
	return 0;
}