Tree POJ - 1741

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

博客https://blog.csdn.net/qq_31759205/article/details/75579558

点分治模板题

#include <cstdio>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;
const int N=0x3f3f3f3f;

struct node
{
    int v;
    int w;
    int next;
};

vector <int> dis;
node edge[20010];
int first[10010],book[10010],sum[10010],maxx[10010];
int n,k,num;

void addedge(int u,int v,int w)
{
    edge[num].v=v;
    edge[num].w=w;
    edge[num].next=first[u];
    first[u]=num++;
}

void getsize(int cur,int fa)
{
    int i,v;
    sum[cur]=1,maxx[cur]=0;
    for(i=first[cur];i!=-1;i=edge[i].next)
    {
        v=edge[i].v;
        if(v!=fa&&book[v]==0)
        {
            getsize(v,cur);
            sum[cur]+=sum[v],maxx[cur]=max(maxx[cur],sum[v]);
        }
    }
}

void getroot(int cur,int fa,int tot,int &minn,int &root)
{
    int i,v;
    maxx[cur]=max(maxx[cur],tot-sum[cur]);
    if(minn>maxx[cur])
    {
        minn=maxx[cur];
        root=cur;
    }
    for(i=first[cur];i!=-1;i=edge[i].next)
    {
        v=edge[i].v;
        if(v!=fa&&book[v]==0) getroot(v,cur,tot,minn,root);
    }
}

void getdis(int cur,int fa,int d)
{
    int i,v,w;
    dis.push_back(d);
    for(i=first[cur];i!=-1;i=edge[i].next)
    {
        v=edge[i].v,w=edge[i].w;
        if(v!=fa&&book[v]==0) getdis(v,cur,d+w);
    }
}

int getnum(int cur,int det)
{
    int res,i,j;
    dis.clear();
    getdis(cur,0,det);
    sort(dis.begin(),dis.end());
    res=0,i=0,j=dis.size()-1;
    while(i<j)
    {
        while(dis[i]+dis[j]>k&&i<j) j--;
        res+=(j-i),i++;
    }
    return res;
}

int dfs(int cur)
{
    int i,v,w,minn,root,res;
    getsize(cur,0);
    minn=N;
    getroot(cur,0,sum[cur],minn,root);
    book[root]=1;
    res=getnum(root,0);
    for(i=first[root];i!=-1;i=edge[i].next)
    {
        v=edge[i].v,w=edge[i].w;
        if(book[v]==0)
        {
            res-=getnum(v,w);
            res+=dfs(v);
        }
    }
    return res;
}

int main()
{
    int i,u,v,w;
    while(scanf("%d%d",&n,&k)!=EOF)
    {
        if(n==0&&k==0) break;
        memset(first,-1,sizeof(first));
        memset(book,0,sizeof(book));
        num=0;
        for(i=1;i<=n-1;i++)
        {
            scanf("%d%d%d",&u,&v,&w);
            addedge(u,v,w);
            addedge(v,u,w);
        }
        printf("%d\n",dfs(1));
    }
    return 0;
}