AcWing252 树 （点分治模板题）

传送门
做一道点分治的裸题
这个题要求一颗树上路径长度小于等于 (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
",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;
}