poj1741 Tree 点分治

入门题，算是对树分治有了初步的理解吧。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<vector>
#define REP(i,a,b) for(int i=a;i<=b;i++)
#define MS0(a) memset(a,0,sizeof(a))using namespace std;typedef long long ll;
const int maxn=20100;
const int INF=1e9+10;int n,k;
int u,v,w;
struct Edge
{int v,w;
};
vector<Edge> G[maxn];
vector<int> d;
bool vis[maxn];void dfs_d(int u,int f,int dep)
{for(int i=0;i<G[u].size();i++){int v=G[u][i].v;if(v==f||vis[v]) continue;dfs_d(v,u,dep+G[u][i].w);d.push_back(dep+G[u][i].w);}
}void dfs_ds(int u,int f,int dep)
{d.push_back(dep);for(int i=0;i<G[u].size();i++){int v=G[u][i].v;if(v==f||vis[v]) continue;dfs_ds(v,u,dep+G[u][i].w);}
}///---get_root
int rt,balance;
int get_root(int u,int f,int sz)
{int cnt=1,balance1=0;for(int i=0;i<G[u].size();i++){int v=G[u][i].v;if(v==f||vis[v]) continue;int tmp=get_root(v,u,sz);cnt+=tmp;balance1=max(balance1,tmp);}balance1=max(balance1,sz-cnt);if(balance1<balance){rt=u;balance=balance1;}return cnt;
}int solve(int u)
{rt=u,balance=INF;int sz=get_root(u,0,n);rt=u,balance=INF;get_root(u,0,sz);u=rt;vis[u]=1;int res=0;d.clear();dfs_d(u,0,0);sort(d.begin(),d.end());for(int i=0;i<d.size();i++){if(d[i]<=k) res++;}for(int l=0,r=(int)d.size()-1;l<r;l++){while(r>l&&d[r]+d[l]>k) r--;res+=r-l;}for(int i=0;i<G[u].size();i++){int v=G[u][i].v;if(vis[v]) continue;d.clear();dfs_ds(v,u,G[u][i].w);sort(d.begin(),d.end());for(int l=0,r=(int)d.size()-1;l<r;l++){while(r>l&&d[r]+d[l]>k) r--;res-=r-l;}}for(int i=0;i<G[u].size();i++){int v=G[u][i].v;if(vis[v]) continue;res+=solve(v);}return res;
}void play()
{freopen("in.txt","w",stdout);n=10000;k=1000000;cout<<n<<" "<<k<<endl;for(int i=1;i<n;i++){cout<<i<<" "<<i+1<<" "<<1<<endl;}cout<<0<<" "<<0<<endl;
}int main()
{//play();return 0;//freopen("in.txt","r",stdin);while(~scanf("%d%d",&n,&k)){if(n==0&&k==0) break;REP(i,1,n) G[i].clear();REP(i,1,n-1){scanf("%d%d%d",&u,&v,&w);G[u].push_back({v,w});G[v].push_back({u,w});}MS0(vis);printf("%d\n",solve(1));}return 0;
}