传送门
做一道点分治的裸题
这个题要求一颗树上路径长度小于等于 \(K\) 的路径的数量。
可以用树状数组维护子树到重心距离在 \([0,K-dis]\) 的节点数量。
但树状数组没法维护 \(0\) 的信息,就同意偏移 \(1\)。
代码
#include <bits/stdc++.h>
#define lowbit(x) ((x)&(-x))
using namespace std;
const int MAXN=1e4+10;
const int MAXM=1e7+10;
int n,K;
int head[MAXN],to[MAXN*2],nxt[MAXN*2],val[MAXN*2],tot;
int rt,siz[MAXN],maxp[MAXN],sum,dead[MAXN];
int ta[MAXM],dis[MAXN],cnt,ans;
void update(int k,int x){
for(int i=k;i<MAXM;i+=lowbit(i)) ta[i]+=x;
}
int ask(int k){
int res=0;
for(int i=k;i>0;i-=lowbit(i)) res+=ta[i];
return res;
}
void add(int u,int v,int w){
to[++tot]=v;nxt[tot]=head[u];val[tot]=w;head[u]=tot;
}
void getrt(int u,int fa){
siz[u]=1;maxp[u]=0;
for(int i=head[u];i;i=nxt[i]){
if(to[i]==fa||dead[to[i]]) continue;
getrt(to[i],u);
siz[u]+=siz[to[i]];
maxp[u]=max(maxp[u],siz[to[i]]);
}
maxp[u]=max(maxp[u],sum-siz[u]);
if(maxp[u]<maxp[rt]) rt=u;
}
void getdis(int u,int fa,int w){
dis[++cnt]=w;
for(int i=head[u];i;i=nxt[i]){
if(to[i]==fa||dead[to[i]]) continue;
getdis(to[i],u,w+val[i]);
}
}
void calc(int u){
update(1,1);
queue<int> que;que.push(1);
for(int i=head[u];i;i=nxt[i]){
if(dead[to[i]]) continue;
cnt=0;
getdis(to[i],u,val[i]);
for(int j=1;j<=cnt;j++) ans+=ask(K-dis[j]);
for(int j=1;j<=cnt;j++) update(dis[j]+1,1),que.push(dis[j]+1);
}
while(!que.empty()) update(que.front(),-1),que.pop();
}
void divide(int u){
dead[u]=1;
calc(u);
for(int i=head[u];i;i=nxt[i]){
if(dead[to[i]]) continue;
maxp[rt=0]=sum=siz[to[i]];
getrt(to[i],0);
getrt(rt,0);
divide(to[i]);
}
}
void solve(){
K++;
for(int i=1,u,v,w;i<n;i++){
scanf("%d%d%d",&u,&v,&w);
u++;v++;
add(u,v,w);add(v,u,w);
}
maxp[0]=sum=n;
getrt(1,0);
getrt(rt,0);
divide(rt);
printf("%d\n",ans);
}
int main(){
while(~scanf("%d%d",&n,&K)&&(n||K)){
tot=ans=0;
memset(head,0,sizeof(head));
memset(dead,0,sizeof(dead));
solve();
}
return 0;
}