POJ 1741.Tree and 洛谷 P4178 Tree-树分治(点分治，容斥版) +二分 模板题-区间点对最短距离<=K的点对数量

POJ 1741.

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 34141		Accepted: 11420

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

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

Sample Output

8

题意就是给你一个带边权的树，求树上最短距离<=K的点对数量。树分治模板题。

代码:

 //树分治-点分治
 #include<cstdio>
 #include<cstring>
 #include<cstdlib>
 #include<cstring>
 #include<algorithm>
 #include<cmath>
 using namespace std;
 typedef long long ll;
 const int inf=1e9+;
 const int maxn=1e5+;

 int head[maxn],tot;
 int root,allnode,ans,n,k;
 int vis[maxn],deep[maxn],dis[maxn],siz[maxn],point[maxn];//deep[0]子节点个数(路径长度)，point为重心节点

 struct node{
     int to,next,val;
 }edge[maxn<<];

 void add(int u,int v,int w)//前向星存图
 {
     edge[tot].to=v;
     edge[tot].next=head[u];
     edge[tot].val=w;
     head[u]=tot++;
 }

 void init()//初始化
 {
     memset(head,-,sizeof head);
     memset(vis,,sizeof vis);
     tot=;
 }

 void get_root(int u,int father)//重心
 {
     siz[u]=;point[u]=;
     for(int i=head[u];~i;i=edge[i].next){
         int v=edge[i].to;
         if(v==father||vis[v]) continue;
         get_root(v,u);//递归得到子树大小
         siz[u]+=siz[v];
         point[u]=max(point[u],siz[v]);//更新u节点的point
     }
     point[u]=max(point[u],allnode-siz[u]);//保存节点size
     if(point[u]<point[root]) root=u;//更新当前子树的重心
 }

 void get_dis(int u,int father)//获取子树所有节点与根的距离
 {
     deep[++deep[]]=dis[u];
     for(int i=head[u];~i;i=edge[i].next){
         int v=edge[i].to;
         if(v==father||vis[v]) continue;
         int w=edge[i].val;
         dis[v]=dis[u]+w;
         get_dis(v,u);
     }
 }

 int cal(int u,int now)
 {
     dis[u]=now;deep[]=;
     get_dis(u,);
     sort(deep+,deep+deep[]+);
     int all=;
     for(int l=,r=deep[];l<r;){//二分
         if(deep[l]+deep[r]<=k){
             all+=r-l;l++;
         }
         else r--;
     }
     return all;
 }

 void solve(int u)//以u为重心进行计算
 {
     ans+=cal(u,);//以当前u为重心的贡献
     vis[u]=;
     for(int i=head[u];~i;i=edge[i].next){
         int v=edge[i].to;
         if(vis[v]) continue;
         ans-=cal(v,edge[i].val);//减去子树的影响
         allnode=siz[v];
         root=;
         get_root(v,u);
         solve(root);
     }
 }

 int main()
 {
     while(~scanf("%d%d",&n,&k)&&n&&k){
         init();
         for(int i=;i<n;i++){
             int u,v,w;
             scanf("%d%d%d",&u,&v,&w);
             add(u,v,w);
             add(v,u,w);
         }
         root=ans=;
         allnode=n;point[]=inf;
         get_root(,);
         solve(root);
         printf("%d\n",ans);
     }
     return ;
 }

洛谷  P4178 Tree

一样的东西，就只是改了一下数据范围和输入格式。

代码:

 //树分治-点分治
 #include<cstdio>
 #include<cstring>
 #include<cstdlib>
 #include<cstring>
 #include<algorithm>
 #include<cmath>
 using namespace std;
 typedef long long ll;
 const int inf=1e9+;
 const int maxn=4e4+;

 int head[maxn],tot;
 int root,allnode,ans,n,k;
 int vis[maxn],deep[maxn],dis[maxn],siz[maxn],point[maxn];//deep[0]子节点个数(路径长度)，point为重心节点

 struct node{
     int to,next,val;
 }edge[maxn<<];

 void add(int u,int v,int w)//前向星存图
 {
     edge[tot].to=v;
     edge[tot].next=head[u];
     edge[tot].val=w;
     head[u]=tot++;
 }

 void init()//初始化
 {
     memset(head,-,sizeof head);
     memset(vis,,sizeof vis);
     tot=;
 }

 void get_root(int u,int father)//重心
 {
     siz[u]=;point[u]=;
     for(int i=head[u];~i;i=edge[i].next){
         int v=edge[i].to;
         if(v==father||vis[v]) continue;
         get_root(v,u);//递归得到子树大小
         siz[u]+=siz[v];
         point[u]=max(point[u],siz[v]);//更新u节点的point
     }
     point[u]=max(point[u],allnode-siz[u]);//保存节点size
     if(point[u]<point[root]) root=u;//更新当前子树的重心
 }

 void get_dis(int u,int father)//获取子树所有节点与根的距离
 {
     deep[++deep[]]=dis[u];
     for(int i=head[u];~i;i=edge[i].next){
         int v=edge[i].to;
         if(v==father||vis[v]) continue;
         int w=edge[i].val;
         dis[v]=dis[u]+w;
         get_dis(v,u);
     }
 }

 int cal(int u,int now)
 {
     dis[u]=now;deep[]=;
     get_dis(u,);
     sort(deep+,deep+deep[]+);
     int all=;
     for(int l=,r=deep[];l<r;){//二分
         if(deep[l]+deep[r]<=k){
             all+=r-l;l++;
         }
         else r--;
     }
     return all;
 }

 void solve(int u)//以u为重心进行计算
 {
     ans+=cal(u,);//以当前u为重心的贡献
     vis[u]=;
     for(int i=head[u];~i;i=edge[i].next){
         int v=edge[i].to;
         if(vis[v]) continue;
         ans-=cal(v,edge[i].val);//减去子树的影响
         allnode=siz[v];
         root=;
         get_root(v,u);
         solve(root);
     }
 }

 int main()
 {
     scanf("%d",&n);
     init();
     for(int i=;i<n;i++){
         int u,v,w;
         scanf("%d%d%d",&u,&v,&w);
         add(u,v,w);
         add(v,u,w);
     }
     scanf("%d",&k);
     root=ans=;
     allnode=n;point[]=inf;
     get_root(,);
     solve(root);
     printf("%d\n",ans);
     return ;
 }