POJ 1741 Tree 【Tree,点分治】

给一棵边带权树,问两点之间的距离小于等于K的点对有多少个。

树上的点 分治,目前还不懂。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <stack>
#include <queue>
#include <vector>
#include <algorithm>
#define mem(x,v) memset(x,v,sizeof(x)) 
#define go(i,a,b)  for (int i = a; i <= b; i++)
#define og(i,a,b)  for (int i = a; i >= b; i--)
using namespace std;
typedef long long LL;
const double EPS = 1e-10;
const int INF = 0x3f3f3f3f;
const int N = 1e4+10;
struct node{
	int val,v;
};
vector<node>g[N];
int n,size,s[N],f[N],root,d[N],K,ans;
vector<int>dep;
bool vis[N];
void getroot(int now, int fa){
	s[now] = 1; f[now] = 0; int u;
	for (int i = 0; i < g[now].size(); i++)
		if ((u = g[now][i].v) != fa && !vis[u]){
			getroot(u,now);
			s[now] += s[u];
			f[now] = max(f[now],s[u]);
		}
	f[now] = max(f[now],size - s[now]);
	if (f[now] < f[root]) root = now;
}
void getdep(int now, int fa){
	int u; dep.push_back(d[now]);
	s[now] = 1;
	for (int i = 0; i < g[now].size(); i++)
		if ((u = g[now][i].v) != fa && !vis[u]){
			d[u] = d[now] + g[now][i].val;
			getdep(u,now);
			s[now] += s[u];
		}

}
int calc(int now, int op){
	dep.clear(); d[now] = op;
	getdep(now,0);
	sort(dep.begin(),dep.end());
	int ret = 0;
	for (int l = 0, r = dep.size() - 1; l < r;)
		if (dep[l] + dep[r] <= K) ret += r-l++; else
		r--;
	return ret;
}
void work(int now){
	int u;
	ans += calc(now,0); vis[now] = 1;
	for(int i = 0; i < g[now].size(); i++)
		if (!vis[u = g[now][i].v]){
			ans -= calc(u,g[now][i].val);
			f[0] = size = s[u];
			getroot(u,root = 0);
			work(root);
		}
}
int main(){
	while(scanf("%d%d",&n,&K) == 2 && n){
		go(i,0,n) g[i].clear();
		mem(vis,0);
		int x,y,z;
		go(i,1,n-1){
			scanf("%d%d%d",&x,&y,&z);
			node e; e.v = y, e.val = z;
			g[x].push_back(e); e.v = x;
			g[y].push_back(e);
		}
		f[0] = size = n;
		getroot(1,root = 0);
		ans = 0;
		work(root);
		printf("%d\n",ans);
	}


	return 0;
}

猜你喜欢

转载自blog.csdn.net/kidsummer/article/details/81704248
今日推荐