1. 程式人生 > 實用技巧 >[A - Distance in Tree] 樹形dp

[A - Distance in Tree] 樹形dp

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\n",#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\n",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\n",u,j,dp[u][j]);
}
ll num[maxn];
void dfs2(int u,int pre) {
//    printf("u=%d pre=%d\n", 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\n",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\n",ans/2);
    return 0;
}