[冬令营day1T3]Tree

题目描述 Description

 给一棵N个节点的无根树，求路径长度=K的简单路径数

输入描述 Input Description

第一行两个正整数N,K

接下来N-1行，每行两个正整数x,y,表示一条边(x,y)

输出描述 Output Description

一行一个整数，答案

样例输入 Sample Input

4 2

1 2

2 3

2 4

样例输出 Sample Output

3

数据范围及提示 Data Size & Hint

对于30% 2<=K,N<=1000

对于100%，2<=K,N<=10^5 

还是一道点分治的题，和POJ1741唯一的区别就是一个是找路径<=k的，一个是找等于k的。两个算法唯一的区别就是在处理dis数组的区别。对于在一个有序表中O(n)的求数对满足两个数的和等于K,我想不到什么比较好的方法，所以导致处理算ans时写的比较丑。

#include<iostream>
#include<cmath>
#include<algorithm>
#include<cstring>
#include<cstdio>
#include<queue>
using namespace std;
typedef long long LL;
#define Pi acos(-1.0)
inline int read()
{
    int x=0,f=1;char c=getchar();
    while(!isdigit(c)){if(c=='-')f=-1;c=getchar();}
    while(isdigit(c)){x=x*10+c-'0';c=getchar();}
    return x*f;
}
const int maxn=100010;
struct Edge
{
    int u,v,next;
    Edge() {}
    Edge(int _1,int _2,int _3) : u(_1),v(_2),next(_3) {}
}e[2*maxn];
int first[maxn],n,k,a,b,size[maxn],masize[maxn],now_size,root,dis[maxn],end;
bool vis[maxn]; 
LL ans;
void addEdge(int i,int a,int b)
{
    e[i]=Edge(a,b,first[a]);
    first[a]=i;
}
void gets(int u,int pa)
{
    size[u]=1;
    masize[u]=0;
    for(int i=first[u];i!=-1;i=e[i].next)
        if(!vis[e[i].v] && e[i].v!=pa)
        {
            gets(e[i].v,u);
            size[u]+=size[e[i].v];
            masize[u]=max(size[e[i].v],masize[u]);
        }
}
void getr(int r,int u,int pa)
{    
    masize[u]=max(masize[u],size[r]-size[u]);
    if(masize[u]<now_size)now_size=masize[u],root=u;
    for(int i=first[u];i!=-1;i=e[i].next)
        if(!vis[e[i].v] && e[i].v!=pa)getr(r,e[i].v,u);
}
void getd(int u,int pa,int d)
{
    dis[end++]=d;
    for(int i=first[u];i!=-1;i=e[i].next)
        if(!vis[e[i].v] && e[i].v!=pa)getd(e[i].v,u,d+1);
}
LL calc(int u,int d)
{
    end=0;
    getd(u,-1,d);
    LL ret=0;
    int l=0,r=end-1,L,R;
    sort(dis,dis+end); 
//    cout<<u<<":";
//    for(int i=0;i<end;i++)cout<<dis[i]<<' ';
//    cout<<endl;
    while(l<r)
    {
        while(dis[l]+dis[r]>k && l<r)r--;
        if(dis[l]+dis[r]==k)
        {
            if(dis[l]==dis[r])
            {
                ret+=((LL)(r-l+1)*(LL)(r-l)/2);
                break;
            }
            else
            {
                L=R=1;
                while(dis[r-1]==dis[r])r--,R++;
                while(dis[l+1]==dis[l])l++,L++;
                ret+=((LL)L*(LL)R);L=R=0;
                r--;l++;
            }
        }
        else l++;
    }
    return ret;
}
void dfs(int u)
{
    now_size=n;
    gets(u,-1);
    getr(u,u,-1);
    ans+=calc(root,0);
    vis[root]=1;
    for(int i=first[root];i!=-1;i=e[i].next)
        if(!vis[e[i].v])
        {
            ans-=calc(e[i].v,1);
            dfs(e[i].v);
        }
    return;
}
int main()
{
    freopen("tree.in","r",stdin);
    freopen("tree.out","w",stdout); 
    memset(first,-1,sizeof(first));
    n=read();k=read();
    for(int i=0;i<n-1;i++)
    {
        a=read();b=read();
        addEdge(i*2,a,b);addEdge(i*2+1,b,a);
    }
    dfs(1);
    printf("%lld
",ans);
    return 0;
}