[A - Distance in Tree] 树形dp

Posted 2020-12-17 echozqn

A - Distance in Tree

CodeForces - 161D

题目大意：树是一个不包含任何圈的连通图。树的两个节点之间的距离是节点之间最短路径的长度（也就是边的长度）。

给定一棵有n个节点的树和一个正整数k，找出距离恰好为k的不同节点对的数量。注意，节点对(v, u)和节点对(u, v)被认为是相同的节点对。

题解：

这个题目可以用点分治写，也可以用换根 (dp) 写。

因为每一条路都是从一个点开始，所以定义 (dp[u][i]) 表示从 (u) 节点出发，路径长度为 (i) 的方案数。

#include <bits/stdc++.h>
#define inf 0x3f3f3f3f
#define inf64 0x3f3f3f3f3f3f3f3f
#define debug(x) printf("debug:%s=%d
",#x,x);
//#define debug(x) cout << #x << ": " << x << endl
# define getchar() (S==T&&(T=(S=BB)+fread(BB,1,1<<20,stdin),S==T)?EOF:*S++)
char BB[1 << 20], *S = BB, *T = BB;
int read(){
    int x=0;
    char c=getchar();
    while (!isdigit(c)) c=getchar();
    while (isdigit(c)) x=(x<<3)+(x<<1)+(c^48),c=getchar();
    return x;
}
using namespace std;
const int maxn = 5e4+10;
typedef long long ll;
int head[maxn<<1],nxt[maxn<<1],to[maxn<<1],cnt;
void add(int u,int v){
    ++cnt,to[cnt]=v,nxt[cnt]=head[u],head[u]=cnt;
    ++cnt,to[cnt]=u,nxt[cnt]=head[v],head[v]=cnt;
}
int n,k;
ll dp[maxn][505];

void dfs1(int u,int pre) {
    dp[u][0] = 1;
//    printf("u=%d pre=%d
",u,pre);
    for (int i = head[u]; i; i = nxt[i]) {
        int v = to[i];
        if (v == pre) continue;
        dfs1(v, u);
        for (int j = 0; j < k; j++) dp[u][j + 1] += dp[v][j];
    }
//    for(int j=0;j<=k;j++) printf("dp[%d][%d]=%I64d
",u,j,dp[u][j]);
}
ll num[maxn];
void dfs2(int u,int pre) {
//    printf("u=%d pre=%d
", u, pre);
    for (int i = head[u]; i; i = nxt[i]) {
        int v = to[i];
        if (v == pre) continue;
        num[0] = 1;
        for (int j = 0; j < k; j++) num[j + 1] = dp[u][j + 1] - dp[v][j];
        for (int j = 0; j < k; j++) dp[v][j + 1] += num[j];
//        for(int j=0;j<=k;j++) printf("dp[%d][%d]=%lld
",v,j,dp[v][j]);
        dfs2(v, u);
    }
}

int main(){
//    freopen("1.in","r",stdin);
    n=read(),k=read();
    for(int i=1;i<n;i++){
        int u = read(),v=read();
        add(u,v);
    }
    dfs1(1,0),dfs2(1,0);
    ll ans=0;
    for(int i=1;i<=n;i++) ans+=dp[i][k];
    printf("%I64d
",ans/2);
    return 0;
}