[省选前题目整理][POJ 1741]Tree(点分治)

题目链接

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

题目大意

求树上使得a到b的最短路小于等于K的无序点对(a,b)个数

思路

经典点分治问题。

如上图，所有符合题意的路径只可能是上面两种形态之一。我们就是要求出所有这样的长度小于等于K的路径，并去掉不合法的路径(重复经过了某条边的路径以及重复计算的路径，实际上点分治是不会出现重复计算的情况的，因为将一个树以重心分割成若干子树后，下一层分治要找的路径分别在各个子树中，不会出现重复)

首先对于当前分治的树，找出它的重心，将重心当成当前分治的树的根节点（注意此时要用一个vis数组标记重心为true，这样在DFS求DFS序、求树的重心等操作时不访问标记过的点，就能保证每次DFS时不会越出当前分治的树的范围！）

然后对这个新的树，从重心(根节点)开始进行DFS，构造出一个DFS序列，并得到每个点的深度，对这个DFS序按照每个点的深度进行升序排序

现在我们要快速地找出所有的点对$(a,b)$,$a,b$是当前子树中的非根节点，并且$depth[a]+depth[b]<=K$

然后维护两个指针$L,R$，初始时分别指向排序后的DFS序的两端，表示$L$和所有的$i$属于$[L+1,R]，并且(L,i)$是符合我们要找的点对的条件的。那么若$depth[L]+depth[R]<=K$，则所有的$(L,i),$i$属于[L+1,R]$显然都是符合$depth[L]+depth[i]<=K$这一条件的，符合上述条件的点对个数就+=$R-L$,$L$加1.否则说明$R$过大，R减1

但是这样做会出现一些不合法情况，如下图，以重心为根节点的子树中，可能出现有两个节点的LCA并不是重心，而是重心下的某个点，这样的点对之间的路径不是最短路

去掉这样的不合法情况的方法也比较简单，即在分治到重心下的某个儿子的子树前，先设这个儿子深度为重心到该儿子的边权，然后以该儿子为根节点DFS这个子树，求出DFS序并排序，和之前做法基本上是一样的，设最终在这个子树中找到的所有满足条件的点对个数为$x$，和之前做法相反，这次是在总答案中减去$x$，因为这些点对都是如上图形态一样的不合法的点对。

代码

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <algorithm>
#include <vector>

#define MAXN 10100

using namespace std;

int n,K,ans=0;

struct edge
{
    int u,v,w,next;
}edges[MAXN*2];

int head[MAXN],nCount=0;

void AddEdge(int U,int V,int W)
{
    edges[++nCount].u=U;
    edges[nCount].v=V;
    edges[nCount].w=W;
    edges[nCount].next=head[U];
    head[U]=nCount;
}

vector<int>seq; //dfs序列
int size[MAXN],vis[MAXN],f[MAXN],root=0; //root=重心,size[i]=子树i大小,f[i]=以i为分治的树的根的话，点i下面最大的儿子子树大小
int sizeOfTree; //当前分治的树的大小
int depth[MAXN];

void getroot(int u,int fa)
{
    size[u]=1;
    f[u]=0;
    for(int p=head[u];p!=-1;p=edges[p].next)
    {
        int v=edges[p].v;
        if(v==fa||vis[v]) continue; //!!!!
        getroot(v,u);
        size[u]+=size[v];
        f[u]=max(f[u],size[v]);
    }
    f[u]=max(f[u],sizeOfTree-size[u]);
    if(f[u]<f[root]) root=u;
}

void getdepth(int u,int fa)
{
    seq.push_back(depth[u]);
    size[u]=1;
    for(int p=head[u];p!=-1;p=edges[p].next)
    {
        int v=edges[p].v;
        if(v==fa||vis[v]) continue;
        depth[v]=depth[u]+edges[p].w;
        getdepth(v,u);
        size[u]+=size[v];
    }
}

int cal(int u,int w) //连接到u的边边权为w,
{
    seq.clear();
    depth[u]=w;
    getdepth(u,0);
    sort(seq.begin(),seq.end());
    int sum=0;
    for(int l=0,r=seq.size()-1;l<r;)
    {
        if(seq[l]+seq[r]<=K)
        {
            sum+=r-l;
            l++;
        }
        else r--;
    }
    return sum;
}

void work(int u)
{
    ans+=cal(u,0);
    vis[u]=true;
    for(int p=head[u];p!=-1;p=edges[p].next)
    {
        int v=edges[p].v;
        if(!vis[v])
        {
            ans-=cal(v,edges[p].w); //删去不合法的路径
            f[0]=sizeOfTree=size[v];
            root=0; //重置树的重心
            getroot(v,0);
            work(root);
        }
    }
}

int main()
{
    while(scanf("%d%d",&n,&K)!=EOF&&!(!n&&!K))
    {
        memset(head,-1,sizeof(head));
        nCount=0;
        root=0;
        ans=0;
        memset(vis,false,sizeof(vis));
        for(int i=1;i<n;i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            AddEdge(u,v,w);
            AddEdge(v,u,w);
        }
        f[0]=sizeOfTree=n; //!!!!!
        getroot(1,0);
        work(root);
        printf("%d\n",ans);
    }
    return 0;
}