点分治-CF161D

题目链接

CF161D-luogu

思路

套用点分治的思想

1. 先处理出每个节点的距离
2. 记录桶中的个数
3. 更新答案
4. 更新桶
5. 下一轮分治清空桶

代码实现

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const int  N = 5e4+10;
int head[N],tot;
int maxp[N],siz[N],rt,sum,dis[N],cnt;
int tmp[N],rec[N],ans,n,k;
bool judge[N];
bool vis[N];
struct edge{
    int v,w,next;
    edge(){}
    edge(int v, int w, int next) : v(v), w(w), next(next) {}
}edges[N<<1];
void add(int u,int v,int w){
    edges[tot]=edge(v,w,head[u]);
    head[u]=tot++;
}
void getrt(int u,int fa){
    maxp[u]=0,siz[u]=1;
    for(int i=head[u];~i;i=edges[i].next){
        int v = edges[i].v;
        if(vis[v]||v==fa)continue;
        getrt(v,u);
        siz[u]+=siz[v];
        if(maxp[u]<siz[v])maxp[u]=siz[v];
    }
    maxp[u]=max(maxp[u],sum-siz[u]);
    if(maxp[rt]>maxp[u])rt=u;
}
void getdis(int u,int fa){
    tmp[cnt++]=dis[u];
    for(int i=head[u];~i;i=edges[i].next){
        int v = edges[i].v;
        if(vis[v]||v==fa)continue;
        dis[v]=dis[u]+edges[i].w;
        getdis(v,u);

    }
}
void solve(int u){
    queue<int>que;
    rec[0]=1;
    for(int i=head[u];~i;i=edges[i].next){
        int v = edges[i].v;
        if(vis[v])continue;
        cnt = 0;
        dis[v]=edges[i].w;
        getdis(v,u);
        for(int j=0;j<cnt;j++){
            if(k-tmp[j]>=0)ans+=rec[k-tmp[j]];
        }
        for(int j=0;j<cnt;j++)rec[tmp[j]]++,que.push(tmp[j]);
    }
    while(!que.empty()){
        rec[que.front()]=0;
        que.pop();
    }
}
void divide(int u){
    vis[u]=1;
    solve(u);
    for(int i=head[u];~i;i=edges[i].next){
        int v = edges[i].v;
        if(vis[v])continue;
        maxp[rt=0]=sum=siz[v];
        getrt(v,0);
        getrt(rt,0);
        divide(rt);
    }
}
int main(){
    memset(head,-1,sizeof(head));
  scanf("%d %d",&n,&k);
  for(int i=1;i<n;i++){
      int u,v,w;
      scanf("%d %d",&u,&v);
      add(u,v,1);
      add(v,u,1);
  }
  maxp[0]=sum=n;
  getrt(1,0);
  getrt(rt,0);
  divide(rt);
  printf("%d\n",ans);
}