POJ 1741 Tree(树上的点分治)

题意：一棵树，边上有权值，求2点之间距离小于等于k的对数。

很经典的点分治问题，理解后思路比较清晰，点对分为3种情况，第1种是在子树中，这就是分治递归的过程。第2种就是子树里的点到根节点距离小于等于k的点，第3种是一个子树里的点到另一个子树的点小于等于k。这个做法就是把点到根节点的距离从小到大进行排序，然后用一个指针j指向末尾，如果第i个跟第j个和大于k，那么i之后的跟j之和肯定也大于k，所以要j--。直到j<=i停止，具体做法可以看代码，这个是o(n)的。然后开始递归子树。还有要开一个标记的数组用于标记那些已经分治过的点，因为需要通过树形dp找这个子树的重心作为根，所以这个子树的递归可能会走到不该走的地方，就是那些已经分治的点。

具体的代码参考了网上的，话说知道原理跟写的确是两码事，看懂别人代码后自己写花了好久的时间调代码。。

AC代码：

#pragma comment(linker, "/STACK:102400000,102400000")
#include<cstdio>
#include<ctype.h>
#include<algorithm>
#include<iostream>
#include<cstring>
#include<vector>
#include<cstdlib>
#include<stack>
#include<queue>
#include<set>
#include<map>
#include<cmath>
#include<ctime>
#include<string.h>
#include<string>
#include<sstream>
#include<bitset>
using namespace std;
#define ll long long
#define ull unsigned long long
#define eps 1e-8
#define NMAX 201000
#define MOD 1000000
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
#define PI acos(-1)
template<class T>
inline void scan_d(T &ret)
{
    char c;
    int flag = 0;
    ret=0;
    while(((c=getchar())<'0'||c>'9')&&c!='-');
    if(c == '-')
    {
        flag = 1;
        c = getchar();
    }
    while(c>='0'&&c<='9') ret=ret*10+(c-'0'),c=getchar();
    if(flag) ret = -ret;
}
const int maxn = 40005;
struct Edge
{
    int v,next,w;
}e[maxn<<1];
int nct,head[maxn],dp[maxn],dmax[maxn];
int a[maxn],all[maxn],t1,t2,ans,k;
bool vis[maxn];
void init()
{
    nct = ans = 0;
    memset(head,-1,sizeof(head));
    memset(vis,0,sizeof(vis));
}

void add(int u, int v, int w)
{
    e[nct].v = v; e[nct].w = w;
    e[nct].next = head[u]; head[u] = nct++;
}

void addedge(int u, int v, int w)
{
    add(u,v,w);
    add(v,u,w);
}

void dfs1(int u,int fa,int &x)
{
    x++;
    dp[u] = 1;
    dmax[u] = 0;
    for(int i = head[u]; i != -1; i = e[i].next)
    {
        int v = e[i].v;
        if(v == fa || vis[v]) continue;
        dfs1(v,u,x);
        dp[u] += dp[v];
        dmax[u] = max(dmax[u],dp[v]);
    }
}
int tmp[2];
void dfs2(int u, int fa, int ge)
{
    int m = max(ge-dp[u], dmax[u]);
    if(tmp[0] == -1 || tmp[1] > m)
    {
        tmp[0] = u;
        tmp[1] = m;
    }
    for(int i = head[u]; i != -1; i = e[i].next)
    {
        int v = e[i].v;
        if(v == fa || vis[v]) continue;
        dfs2(v,u,ge);
    }
}

int getbarycenter(int u)
{
    int ge=0;
    dfs1(u,-1,ge);
    tmp[0] = -1;
    dfs2(u,-1,ge);
    return tmp[0];
}

void cal(int u, int fa, int add)
{
    a[t1++] = add;
    all[t2++] = add;
    for(int i = head[u]; i != -1; i = e[i].next)
    {
        int v = e[i].v;
        if(v == fa || vis[v]) continue;
        cal(v,u,add+e[i].w);
    }
}

int work(int *t, int x)
{
    int tmp = 0,j = x-1;
    sort(t,t+x);
    for(int i = 0; i < x; i++)
    {
        while(t[i] + t[j] > k) j--;
        if(j <= i) break;
        tmp += j-i;
    }
    return tmp;
}

void solve(int u)
{
//    if(u == -1) return;
    u = getbarycenter(u);
    all[0] = 0;
    t2 = 1;
    for(int i = head[u]; i != -1; i = e[i].next)
    {
        int v = e[i].v;
        if(vis[v]) continue;
        t1 = 0;
        cal(v,u,e[i].w);
//        cout<<"kao:"<<a[0]<<" "<<v<<endl;
//        for(int i = 0; i < t1; i++)
//            cout<<a[i]<<" ";
//        cout<<endl;
        ans -= work(a,t1);
    }
    ans += work(all,t2);
    vis[u] = 1;
    for(int i = head[u]; i != -1; i = e[i].next)
    {
        int v = e[i].v;
        if(vis[v]) continue;
        solve(v);
    }
}

int main()
{
#ifdef GLQ
    freopen("input.txt","r",stdin);
//    freopen("o3.txt","w",stdout);
#endif
    int n;
    while(~scanf("%d%d",&n,&k) && n+k)
    {
        init();
        for(int i = 0; i < n-1; i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            addedge(u,v,w);
        }
        solve(1);
        printf("%d\n",ans);
    }
    return 0;
}