2017.7.18 树上距离=k 思考记录

来源:互联网 发布:java防止sql注入转义 编辑:程序博客网 时间:2024/06/06 03:12

利用树分治的原理,再加上扫描线的o(nm)更新,可以求出树上任意两点距离==k的点对个数



#include<iostream>#include<cstdio>#include<cstring>#include<algorithm>using namespace std;#define N 20005#define inf 1e9+7int hou[N],xia[N],zhong[N],v[N],tot,max1,sz[N],i,rt,dis[N],q[N],n,m,w[N],x,y,z,ans[N];bool vis[N];void jia(int x,int y,int z){++tot,hou[tot]=xia[x],zhong[tot]=y,xia[x]=tot,v[tot]=z;//cout<<tot<<" "<<x<<" "<<y<<endl;}void jian(int x,int y,int z){jia(x,y,z);jia(y,x,z);}int zzx(int o,int fu){sz[o]=1;int maxx=0;for(int i=xia[o];i!=-1;i=hou[i]){int nd=zhong[i];if(nd==fu||vis[nd])continue; zzx(nd,o); sz[o]+=sz[nd];     maxx=max(maxx,sz[nd]);} maxx=max(maxx,tot-maxx);if(maxx<max1){max1=maxx,rt=o;}}void dfs(int o,int fu,int d){dis[o]=d;q[++q[0]]=d;for(int i=xia[o];i!=-1;i=hou[i]){int nd=zhong[i];//cout<<o<<" "<<nd<<endl;if(nd==fu||vis[nd])continue;dfs(nd,o,d+v[i]); } }void tj1(int o,int fu){q[0]=0;dfs(o,fu,0);    sort(q+1,q+1+q[0]);    for(int i=1;i<=m;i++)    {    int l=1;    int r=q[0];    while(l<r)    {    if(q[l]+q[r]==w[i]){ans[i]++;if(q[l]+q[r-1]==w[i])--r;else ++l;}    else if(q[l]+q[r]>w[i])--r;    else ++l;}}}void tj2(int o,int fu){q[0]=0;dfs(o,fu,dis[o]-dis[fu]);  sort(q+1,q+1+q[0]);    for(int i=1;i<=m;i++)    {    int l=1;    int r=q[0];    while(l<r)    {    if(q[l]+q[r]==w[i]){ans[i]--;if(q[l]+q[r-1]==w[i])--r;else ++l;}    else if(q[l]+q[r]>w[i])++l;    else --r;}}}void work(int o,int fu){vis[o]=1;   tj1(o,fu);for(int i=xia[o];i!=-1;i=hou[i]){int nd=zhong[i];if(nd==fu)continue;tj2(nd,o);       tot=sz[nd];       max1=inf;zzx(nd,o); work(rt,o);}}int main(){memset(xia,-1,sizeof(xia));scanf("%d%d",&n,&m);for(i=1;i<n;i++){scanf("%d%d%d",&x,&y,&z);jian(x,y,z);}for(i=1;i<=m;i++){scanf("%d",&w[i]);}tot=n;max1=inf;zzx(1,0);work(rt,0);for(i=1;i<=m;i++){printf("%d ",ans[i]);}}