本题如果k的范围较小的话,可以使用树状数组记录答案,但是因为很大
考虑使用双指针+容斥原理。
也就是直接算整个子树的答案,之后再在枚举儿子节点的时候,把加上u-v这条边的合法答案全部清除,这样就做到了不重不漏
#include<bits/stdc++.h> using namespace std; typedef long long ll; typedef pair<int,int> pll; const int N=1e5+10; int h[N],ne[N],e[N],w[N],idx; int n,K; int root,cnt; int d[N],vis[N]; int sz[N]; void add(int a,int b,int c){ e[idx]=b,ne[idx]=h[a],w[idx]=c,h[a]=idx++; } void dfs_rt(int u,int fa,int tot){ sz[u]=1; int i; int n=0; for(i=h[u];i!=-1;i=ne[i]){ int j=e[i]; if(j==fa||vis[j]) continue; dfs_rt(j,u,tot); sz[u]+=sz[j]; n=max(n,sz[j]); } n=max(n,tot-sz[u]); if(n*2<=tot) root=u; } void dfs_dis(int u,int fa,int x){ d[++cnt]=x; int i; for(i=h[u];i!=-1;i=ne[i]){ int j=e[i]; if(j==fa||vis[j]) continue; dfs_dis(j,u,x+w[i]); } } int get(int u,int x){ cnt=0; dfs_dis(u,-1,x); int l=1,r=cnt; int ans=0; sort(d+1,d+1+cnt); while(l<r){ while(d[l]+d[r]>K&&l<r) r--; ans+=r-l; l++; } return ans; } void dfs_sz(int u,int fa){ sz[u]=1; for(int i=h[u];i!=-1;i=ne[i]){ int j=e[i]; if(j==fa||vis[j]) continue; dfs_sz(j,u); sz[u]+=sz[j]; } } int work(int u,int tot){ dfs_rt(u,-1,tot); u=root; dfs_sz(u,-1); vis[u]=1; int i; int res=get(u,0); //cout<<res<<endl; for(i=h[u];i!=-1;i=ne[i]){ int j=e[i]; if(!vis[j]){ res-=get(j,w[i]); res+=work(j,sz[j]); } } return res; } int main(){ ios::sync_with_stdio(false); while(cin>>n>>K){ if(n==0&&K==0) break; int i; memset(h,-1,sizeof h); memset(vis,0,sizeof vis); for(i=1;i<n;i++){ int a,b,c; cin>>a>>b>>c; a++,b++; add(a,b,c); add(b,a,c); } cout<<work(1,n)<<endl; } }View Code