【树的点分治模版】POJ 1741 Tree

http://poj.org/problem?id=1741

从来没打过版，打了一个

五个变量&数组：cnt、mx、root、size、flag

cnt、mx只在getroot中更新

divide函数很简单，只做分治的工作

cal函数负责计算，用getdis预处理

注意cal函数里的memset，有可能会影响复杂度

我就掉这坑里了，把memset去掉之后快了8倍！

注意不用写暴力而对拍的技巧

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <iostream>
#include <cmath>
#include <algorithm>
#include <vector>
using namespace std;
#define rep(i,l,r) for(int i=(l),_=(r);i<=_;i++)
#define per(i,r,l) for(int i=(r),_=(l);i>=_;i--)
#define MS(arr,x) memset(arr,x,sizeof(arr))
#define INE(i,u) for(int i=head[u];~i;i=e[i].next)
#define LL long long
inline const int read()
{int r=0,k=1;char c=getchar();for(;c<'0'||c>'9';c=getchar())if(c=='-')k=-1;
for(;c>='0'&&c<='9';c=getchar())r=r*10+c-'0';return k*r;}

const int N=10010;
int n,m;
struct edge{int v,w,next;}e[N*2];
int head[N],k;
int cnt[N],mx[N],root,size;
bool flag[N];
int ans;
int dis[N];
vector<int>d;

void adde(int u,int v,int w){e[k]=(edge){v,w,head[u]};head[u]=k++;}
inline void MAX(int &a,int b){if(a<b)a=b;}
void getroot(int u,int fa)
{
	cnt[u]=1; mx[u]=0;
	INE(i,u)
	{
		int v=e[i].v; if(v == fa || flag[v]) continue;
		getroot(v,u);
		cnt[u]+=cnt[v];
		MAX(mx[u],cnt[v]);
	}
	MAX(mx[u],size-cnt[u]);
	if(mx[u] < mx[root]) root=u;
}
void getdis(int u,int fa)
{
	d.push_back(dis[u]);
	INE(i,u)
	{
		int v=e[i].v; if(v == fa || flag[v]) continue;
		dis[v]=dis[u]+e[i].w;
		getdis(v,u);
	}
}
int cal(int u,int shift)
{
	d.clear(); dis[u]=0; //MS(dis,0);！！！！！！！！！！！！！！！ 
	getdis(u,0);
	sort(d.begin(),d.end());
	int res=0;
	for(int l=0,r=d.size()-1;l<r;)
	    if(d[l] + d[r] <= m - shift) res+=r-l++;
	    else r--;
	return res;
}
void divide(int u)
{
	ans+=cal(u,0);
	flag[u]=1;
	INE(i,u)
	{
		int v=e[i].v; if(flag[v]) continue;
		ans-=cal(v,2*e[i].w);
		size=mx[root=0]=cnt[v];
		getroot(v,u);
		divide(root);
	}
}

void input()
{
    MS(head,-1); k=0; MS(flag,0);
    rep(i,2,n)
    {
    	int u=read(),v=read(),w=read();
    	adde(u,v,w); adde(v,u,w);
    }
}
void solve()
{
	size=mx[root=0]=n;
	getroot(1,-1);
	ans=0;
	divide(root);
	printf("%d\n",ans);
}

int main()
{
    freopen("_.in","r",stdin); freopen("_.out","w",stdout);
    while(scanf("%d%d",&n,&m),n+m)
    input(),solve();
    return 0;
}