POJ 1741 Tree（点分治）

http://poj.org/problem?id=1741

题意：

给出一棵树，求出树上满足两点权值之和不大于k的点对数。

思路：

直接暴力就是$O(n^2)$，显然不行。这里有一篇论文可以推荐大家看一下。

因为一条路径要么过根结点，要么不过，即在一棵子树中，然后我们可以用分治法。每次确定一个根u，u有几棵子树，就将u的子树分成了几部分，每一部分就一直递归下去。

我们设dis[i]为i到根结点的距离，于是我们可以得到以下的算法：$dis[i]+dis[j]<=k的(i,j)对数-dis[i]+dis[j]<=k的(i,j)对数且(i,j)在一棵子树中$。这个部分其实不难算，我们只需要先把所有子节点的dis值求出，然后排个序，利用尺取法可以快速求出个数。

还有一点要注意的是，根结点的选择很重要，这里需要选择重心，因为删去重心后最大子树的节点数最小。

#include<iostream>
#include<algorithm>
#include<cstring>
#include<cstdio>
#include<vector>
#include<stack>
#include<queue>
#include<cmath>
#include<map>
#include<set>
using namespace std;
typedef long long ll;
typedef pair<int,int> pll;
const int INF = 0x3f3f3f3f;
const int maxn=10000+5;

int n,k,num,root,sz,ans,mi;
int vis[maxn], son[maxn], dis[maxn], mx[maxn];
vector<pll> G[maxn];


void getroot(int u, int fa, int n) //计算重心
{
    son[u]=0;
    int balance=0;
    for(int i=0;i<G[u].size();i++)
    {
        int v=G[u][i].first;
        if(v==fa || vis[v])  continue;
        getroot(v,u,n);
        son[u]+=son[v]+1;
        balance=max(balance,son[v]+1);
    }
    balance=max(balance,n-son[u]-1);
    if(balance<sz)
    {
        sz=balance;
        root=u;
    }
}


void getdis(int u, int fa, int d)  //求子节点到根结点的距离
{
    dis[num++]=d;
    for(int i=0;i<G[u].size();i++)
    {
        int v=G[u][i].first;
        int w=G[u][i].second;
        if(v==fa || vis[v])  continue;
        getdis(v,u,d+w);
    }
}

int getnum()
{
    int ret=0;
    sort(dis,dis+num);
    int l=0,r=num-1; 
    while(l<r)  //尺取法快速求对数
    {
        while(dis[l]+dis[r]>k && l<r)  r--;
        ret+=r-l;
        l++;
    }
    return ret;
}

void solve(int u, int n)
{
    sz=INF;
    getroot(u,-1,n);
    num=0;
    getdis(root,-1,0);
    ans+=getnum();
    vis[root]=1;
    int tmp=root;
    for(int i=0;i<G[tmp].size();i++)
    {
        int v=G[tmp][i].first;
        int w=G[tmp][i].second;
        if(!vis[v])
        {
            num=0;
            getdis(v,-1,w);  //减去在同一棵子树的情况
            ans-=getnum();
            solve(v,son[v]+1);
        }
    }
}

int main()
{
    //freopen("in.txt","r",stdin);
    while(~scanf("%d%d",&n,&k))
    {
        if(!n && !k)  break;
        memset(vis,0,sizeof(vis));
        for(int i=1;i<=n;i++)  G[i].clear();
        for(int i=1;i<n;i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            G[u].push_back(make_pair(v,w));
            G[v].push_back(make_pair(u,w));
        }
        ans=0;
        solve(1,n);
        printf("%d
",ans);
    }
    return 0;
}