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;
}