P4178 Tree

题目描述

给你一棵TREE,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K

输入输出格式

输入格式：

N（n<=40000） 接下来n-1行边描述管道，按照题目中写的输入 接下来是k

输出格式：

一行，有多少对点之间的距离小于等于k

输入输出样例

输入样例#1： 复制

7
1 6 13 
6 3 9 
3 5 7 
4 1 3 
2 4 20 
4 7 2 
10

输出样例#1： 复制

5

//Pro:P4178 Tree

//点分治

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;

inline int read()
{
    char c=getchar();int num=0;
    for(;!isdigit(c);c=getchar());
    for(;isdigit(c);c=getchar())
        num=num*10+c-'0';
    return num;
}

const int N=4e4+5;
const int INF=599518803;

int n,k,ans;

int head[N],num_edge;
struct Edge
{
    int v,w,nxt;
}edge[N<<1];

inline void add_edge(int u,int v,int w)
{
    edge[++num_edge].v=v;
    edge[num_edge].w=w;
    edge[num_edge].nxt=head[u];
    head[u]=num_edge;
}

bool vis[N];

int root,Siz;
int siz[N],mxson[N];
void getroot(int u,int fa)
{
    siz[u]=1,mxson[u]=0;
    for(int i=head[u],v;i;i=edge[i].nxt)
    {
        v=edge[i].v;
        if(vis[v]||v==fa)
            continue;
        getroot(v,u);
        siz[u]+=siz[v];
        mxson[u]=max(mxson[u],siz[v]);
    }
    mxson[u]=max(mxson[u],Siz-siz[u]);
    if(mxson[u]<mxson[root])
        root=u;
}

int dis[N],dep[N],cnt;
void getdis(int u,int fa)
{
    dis[++cnt]=dep[u];
    for(int i=head[u],v;i;i=edge[i].nxt)
    {
        v=edge[i].v;
        if(v==fa||vis[v])
            continue;
        dep[v]=dep[u]+edge[i].w;
        getdis(v,u);
    }
}

int solve(int u,int dist)
{
    cnt=0,dep[u]=dist;
    getdis(u,u);
    sort(dis+1,dis+cnt+1);
    int l=1,r=cnt,res=0;
    while(l<r)
    {
        if(dis[l]+dis[r]<=k)
        {
            res+=r-l;
            ++l;
        }
        else
            --r;
    }
    return res;
}

void divide(int u)
{
    ans+=solve(u,0);
    vis[u]=1;
    for(int i=head[u],v;i;i=edge[i].nxt)
    {
        v=edge[i].v;
        if(vis[v])
            continue;
        ans-=solve(v,edge[i].w);
        Siz=siz[v],root=0;
        getroot(v,v);
        divide(root);
    }
}

int main()
{
    n=read();
    for(int i=1,a,b,c;i<n;++i)
    {
        a=read(),b=read(),c=read();
        add_edge(a,b,c);
        add_edge(b,a,c);
    }
    k=read();
    Siz=n,mxson[0]=INF;
    getroot(1,1);
    divide(root);
    printf("%d",ans);
    return 0;
}