poj 1741 Tree 树分治

题目连接：点击打开链接

题意：给你一棵带边权的树，求树上两点距离小于等于k的对数。点数n<=10000，k<=1000

解题思路就是点分治。如果不清楚树分治的思想可以查看漆子超关于树分治的论文。

然后具体思路就是，将两点之间的路径分为两类，一类是经过某个节点，一类是在某个节点的子树里面。然后为了保证时间复杂度能在logn，每次必须找树的重心。

说说我遇到的问题吧！我一直TLE，主要是没有彻底的应用分治的思想，在每次算一棵子树的时候要统计所有节点的dep[]，而不是统计f[i]表示到该子树根节点为 i 的节点的数目。这时候每次分治是一个固定的常数max_k，所以一直TLE。

#include<stdio.h>
#include<iostream>
#include<string.h>
#include<math.h>
#include<algorithm>
#include<vector>
#include<map>
#include<set>
#include<queue>
#include<string>
#define ll long long
#define db double
using namespace std;

const int N = 10005;
const int INF= 1000000000;
const int K = 1005;

int max_k;

int head[N],to[2*N],cost[2*N],next[2*N];
int nedge;
bool vis[N];

void init()
{
    memset(head,-1,sizeof(head));
    nedge=0;
}

void add(int u,int v,int c)
{
    to[nedge]=v,cost[nedge]=c,next[nedge]=head[u],head[u]=nedge++;
}

ll ans;

int root;
int tol[N],max_node;

void get_root(int k,int fa,int &n)
{
    tol[k]=0;
    int mm=0;
    for(int i=head[k];i>=0;i=next[i])
    {
        if(to[i]==fa||vis[to[i]]) continue;
        get_root(to[i],k,n);
        tol[k]+=tol[to[i]];
        mm=max(mm,tol[to[i]]);
    }
    tol[k]++;
    mm=max(mm,n-tol[k]);
    if(mm<max_node)
    {
        max_node=mm;
        root=k;
    }
}

int dep[N],cnt;

void get_dep(int k,int fa,int d)
{
    dep[cnt++]=d;
    for(int i=head[k];i>=0;i=next[i])
    {
        if(to[i]!=fa&&!vis[to[i]])
        {
            get_dep(to[i],k,d+cost[i]);
        }
    }
}

void work(int k,int fa,int n)
{
    max_node=INF;
    get_root(k,fa,n);
    int rt=root;
    vis[rt]=true;
    for(int i=head[rt];i>=0;i=next[i])
    {
        if(to[i]==fa||vis[to[i]]) continue;
        if(tol[to[i]]>tol[rt]) work(to[i],rt,n-tol[rt]);
        else work(to[i],rt,tol[to[i]]);
    }
    cnt=0;
    get_dep(rt,rt,0);
    sort(dep,dep+cnt);
    int l=0,r=cnt-1;
    while(l<r)
    {
        if(dep[l]+dep[r]<=max_k) ans+=r-l,l++;
        else r--;
    }
    for(int i=head[rt];i>=0;i=next[i])
    {
        if(to[i]!=fa&&!vis[to[i]]){
            cnt=0;
            get_dep(to[i],rt,cost[i]);
            sort(dep,dep+cnt);
            l=0,r=cnt-1;
            while(l<r)
            {
                if(dep[l]+dep[r]<=max_k) ans-=r-l,l++;
                else r--;
            }
        }
    }
    vis[rt]=false;
}

int main()
{
#ifdef PKWV
    freopen("in.in","r",stdin);
#endif // PKWV
    int n;
    while(scanf("%d%d",&n,&max_k)&&n+max_k)
    {
        init();
        for(int i=1;i<n;i++)
        {
            int a,b,c;
            scanf("%d%d%d",&a,&b,&c);
            add(a,b,c),add(b,a,c);
        }
        ans=0;
        work(1,-1,n);
        cout<<ans<<endl;
    }
    return 0;
}