poj1741 Tree[点分治]

2018-01-12

Description

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.

Input

The input contains several test cases. The first line of each test case contains two integers $n, k$ . ( $n<=10000$ ) The following $n-1$ lines each contains three integers $u,v,l$ , which means there is an edge between node $u$ and $v$ of length $l$ .
The last test case is followed by two zeros.

Output

For each test case output the answer on a single line.

Sample Input

Sample Output

Solution

模板，

和洛谷的模板差不多，不过luogu的是询问树上距离为 $k$ 的点对是否存在（因为范围不大直接开桶暴力就好）。这道题是给一颗树，求树上距离 $<=k$ 的点对有多少个。

每个点去统答案，然后去掉两端都在同一颗子树的情况（luogu那题可以直接标记该树属于哪个子树）

#include<cmath>
#include<cstdio>
#include<string>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
int qmax(int &x,int y) {if (x<y) x=y;}
int qmin(int &x,int y) {if (x>y) x=y;}
int read()
{
    char s;int k=0,base=1;
    while((s=getchar())!='-'&&s!=EOF&&!(isdigit(s)));if(s=='-')base=-1,s=getchar();
    while(isdigit(s)){k=k*10+(s^'0');s=getchar();}
    return k*base;
}
const int maxn=1e4+100;
int n,k,id,to[maxn<<1],ne[maxn<<1],w[maxn<<1];
int vis[maxn],po[maxn],X,Y,Z,root,SZ;
int deep[maxn],d0,size[maxn],son[maxn];
int a[maxn],ans; 
inline void add(int x,int y,int z)
{
    id++;to[id]=y;w[id]=z;ne[id]=po[x];po[x]=id;
}
void getroot(int x,int fa)//找重心
{
    son[x]=0;size[x]=1;
    for (int i=po[x];i;i=ne[i])
    {
        if (!vis[to[i]]&&to[i]!=fa)//被打了标记相当于删了
        {
            getroot(to[i],x);
            size[x]+=size[to[i]];
            qmax(son[x],size[to[i]]);
        }
    }
    son[x]=max(son[x],SZ-size[x]);//上面那部分也要统上去
    if (son[x]<son[root]||root==0) root=x;
}
void getdeep(int x,int fa)//距离
{
    a[++d0]=deep[x];
    for (int i=po[x];i;i=ne[i])
    {
        if (!vis[to[i]]&&to[i]!=fa)
        {
            deep[to[i]]=deep[x]+w[i];
            getdeep(to[i],x);
        }
    }
}
int cal(int x,int y)//统计x->y的ans
{
    deep[x]=y;d0=0;int ss=0;
    getdeep(x,0);
    sort(a+1,a+d0+1);
    int l=1,r=d0;
    while (l<r)
    {
        if (a[r]+a[l]<=k) ss+=(r-l),l++; else r--;
    }
    return ss;
} 
void solve(int x)
{
    ans+=cal(x,0);vis[x]=1;
    for (int i=po[x];i;i=ne[i])
    {
        if (!vis[to[i]])
        {
            ans-=cal(to[i],w[i]);//去不合法的
            root=0;SZ=size[to[i]];
            getroot(to[i],0);//找子树重心
            solve(root);//处理子树
        }
    }
}
int main()
{
    n=read();k=read();
    while (n&&k)
    {
        ans=0;
        memset(vis,0,sizeof(vis));
        id=0;memset(po,0,sizeof(po));
        memset(size,0,sizeof(size));d0=0;
        for (int i=1;i<n;i++)
        {
            X=read();Y=read();Z=read();
            add(X,Y,Z);add(Y,X,Z);
        }
        SZ=n;
        getroot(1,0);//找重心
        solve(root);
        printf("%d\n",ans);
        n=read();k=read();
    }
}