POJ - 1741,点分治

A - Tree

Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.
题意：给你一棵树，问其中有多少个点对路径长度小于等于k
思路：点分治板题，每次枚举两条到重心的路径时可以先记录一个路径出现次数的前缀和，从而可以O(n)枚举所有距离相加<=k的路径条数

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define MAXN 10010
#define INF 0x3f3f3f3f
using namespace std;
int head[MAXN],tot;
struct edge
{
    int v,w,nxt;
}edg[MAXN << 1];
inline void addedg(int u,int v,int w)
{
    edg[tot].v = v;
    edg[tot].w = w;
    edg[tot].nxt = head[u];
    head[u] = tot++;
}
int n,root,ms,mson[MAXN],sz[MAXN],Size;
bool vis[MAXN];
//root用于标记重心,ms表示树的重心的最大子树的大小,mson[i]记录以i为根最大子树的大小
//sz[i]记录以i为根子树的大小,Size表示当前整棵树的大小,vis[i]表示当前节点是否被分治过
void getroot(int u,int f)//获得重心
{
    sz[u] = 1,mson[u] = 0;
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v] || v == f) continue;//剔除已经被分治过的点
        getroot(v,u);
        sz[u] += sz[v];
        if(sz[v] > mson[u]) mson[u] = sz[v];
    }
    if(Size - sz[u] > mson[u]) mson[u] = Size-sz[u];//把u看作根节点时u的父亲那一部分也算作子树
    if(ms > mson[u]) ms = mson[u],root = u;//更新重心
}
int k,ans;
int dis[MAXN],cnt;//dis记录所有节点到重心的距离
struct node
{
    int d,num;
}nod[MAXN];//记录距离d出现的次数num
int cc;//不同距离数
void getdis(int u,int f,int d)//获得到目标点的距离
{
    dis[++cnt] = d;
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v] || v == f) continue;
        getdis(v,u,d + edg[i].w);
    }
}
void cal(int u,int d,int tp)//u表示getdis的起点,d表示u到目标点的距离,tp表示这一次统计出来的答案是合理的还是不合理的
{
    cnt = 0;
    getdis(u,0,d);//算出树中的点到目标点的距离
    sort(dis+1,dis+cnt+1);
    nod[0].num = 0;
    cc = 1,nod[cc].d = dis[1],nod[cc].num = 1;
    for(int i = 2;i <= cnt;++i)
        if(dis[i] != dis[i-1]) nod[++cc].d = dis[i],nod[cc].num = nod[cc-1].num+1;
        else ++nod[cc].num;
    int l = 1,r = cc;
    while(l <= r)
    {
        if(nod[l].d + nod[r].d <= k)
        {
            if(l < r)
                ans += (nod[l].num-nod[l-1].num)*(nod[r].num-nod[l].num) * tp;
            ans += (nod[l].num-nod[l-1].num)*(nod[l].num-nod[l-1].num-1)/2*tp;
            ++l;
        }
        else
            --r;
    }
}
void solve(int u,int ssize)//ssize是当前这棵子树的大小
{
    vis[u] = true;//代码保证每次进来的u都必定是当前这棵树的重心,我们将vis[u]标记为true,表示u点被分治过
    cal(u,0,1);//计算这棵树以u为重心的所有组合,但包括了共用同一条边的情况
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v]) continue;
        cal(v,edg[i].w,-1);//将共用一条边的不合法情况去除
        ms = INF;//记得每次都要初始化
        Size = sz[v] < sz[u]?sz[v]:(ssize-sz[u]);//因为v实际上可能是u的父亲,故sz需相减
        getroot(v,v);//求出以v为根节点的子树重心
        solve(root,Size);
    }
}
inline void init()
{
    tot = 0,ms = INF,Size = n;
    memset(head,-1,sizeof(int)*(n+1));
    memset(vis,false,sizeof(bool)*(n+1));
    ans = 0;
}
int main()
{
    while(~scanf("%d%d",&n,&k) && (n || k))
    {
        init();
        int u,v,w;
        for(int i = 1;i < n;++i)
        {
            scanf("%d%d%d",&u,&v,&w);
            addedg(u,v,w),addedg(v,u,w);
        }
        getroot(1,1);
        solve(root,Size);
        printf("%d\n",ans);
    }
    return 0;
}