POJ 1741 树分治入门

#include<iostream>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
const int INF=0x3f3f3f3f;
const int MAX=1e4+5;
//点分治 CurSize-树当前大小 CurRoot-当前树的根
//Son[i]节点i的子节点的个数 D[i]节点i到CurRoot的距离 F[i]节点i最大子树的大小
int N,K,CurSize,CurRoot,Ans,Son[MAX],D[MAX],F[MAX];
bool Used[MAX]; //Used[i]节点i是否计算完成
vector<pair<int,int> > G[MAX]; //Graph
vector<int> Depth;  //把计算后的D[i]保存进来
void GetRoot(int s,int pre)
{
Son[s]=0,F[s]=0;
for (int i=0; i<(int)G[s].size(); i++)
{
if (G[s][i].first==pre||Used[G[s][i].first])
continue;
GetRoot(G[s][i].first,s);
Son[s]+=Son[G[s][i].first]+1;
F[s]=max(F[s],Son[G[s][i].first]+1);
}
F[s]=max(F[s],CurSize-Son[s]-1);
if (F[s]<F[CurRoot])
CurRoot=s;
}
void GetDepth(int s,int pre)
{
Depth.push_back(D[s]);
Son[s]=0;
for (int i=0; i<(int)G[s].size(); i++)
{
if (G[s][i].first==pre||Used[G[s][i].first])
continue;
D[G[s][i].first]=D[s]+G[s][i].second;
GetDepth(G[s][i].first,s);
Son[s]+=Son[G[s][i].first]+1;
}
}
int Calc(int s,int init)
{
Depth.clear();
D[s]=init;
GetDepth(s,0);
sort(Depth.begin(),Depth.end());
int res=0;
for (int l=0,r=Depth.size()-1; l<r;)
if (Depth[l]+Depth[r]<=K)
res+=r-l++;
else
r--;
return res;
}
void Solve(int s)
{
Ans+=Calc(s,0); //求出所有D[i]+D[j]<=K的路径的个数
Used[s]=true;
for (int i=0; i<(int)G[s].size(); i++)
{
if (!Used[G[s][i].first])
{
Ans-=Calc(G[s][i].first,G[s][i].second); //再减去i和j在同一子树的情况
F[0]=CurSize=Son[G[s][i].first]+1;
GetRoot(G[s][i].first,CurRoot=0);
Solve(CurRoot);
}
}
}
int main()
{
cin.sync_with_stdio(false);
while(cin>>N>>K)
{
if (N==0&&K==0)
break;
int a,b,c;
memset(Used,false,sizeof(Used));
for (int i=1; i<=N; i++)
G[i].clear();
for (int i=1; i<N; i++)
{
cin>>a>>b>>c;
G[a].push_back(make_pair(b,c));
G[b].push_back(make_pair(a,c));
}
F[0]=CurSize=N;
GetRoot(1,CurRoot=0);
Ans=0;
Solve(CurRoot);
cout<<Ans<<endl;
}
return 0;
}