cf161d 求距离为k的点对(点分治，树形dp)

点分治裸题，但是用树形dp也能做

/*
dp[u][k]表示在u下距离k的点数量 
*/
#include<bits/stdc++.h>
using namespace std;
struct Edge{int to,nxt;}edge[100005];
int head[50005],tot,n,k;
long long dp[50005][505],ans;
void init(){
    memset(head,-1,sizeof head);
    tot=0;
}
void addedge(int u,int v){
    edge[tot].to=v;edge[tot].nxt=head[u];head[u]=tot++;
}
void dfs(int u,int pre){ 
    dp[u][0]++;
    for(int i=head[u];i!=-1;i=edge[i].nxt){
        int v=edge[i].to;
        if(v!=pre){
            dfs(v,u);
            for(int j=0;j<k;j++)
                dp[u][j+1]+=dp[v][j];
        }
    }

    for(int i=head[u];i!=-1;i=edge[i].nxt){
        int v=edge[i].to;
        if(v==pre)continue;
        for(int j=0;j<k;j++){
            ans+=dp[v][j]*dp[u][k-j-1];//距离v为j的点数*距离u为k-j-1的点数 
            if(k-j-1>0)ans-=dp[v][j]*dp[v][k-j-2];//距离v为j的点数*距离v为k-j-2的点数 
        }
    }
    ans+=dp[u][k];
}
int main(){
    init();
    int u,v;
    cin>>n>>k;
    for(int i=1;i<n;i++){
        cin>>u>>v;
        addedge(u,v);addedge(v,u);
    }    
     
    dfs(1,0);
    printf("%d\n",ans/2);
}

 

转载于:https://www.cnblogs.com/zsben991126/p/10333323.html