CFR 161 D. Distance in Tree ( Tree DP )

Problem - 161D - Codeforces

題意：
給一棵樹，求滿足距離為 K 的點對數。
解法：O( N * K )
cdis[ u ][ i ] : count of nodes v, such that v is part of the subtree of u, and is exactly i edges aways from u
cdis[ u ][ i ] = SIGMA{ cdis[ v ][ i - 1 ] | v is part of the subtree of u }, if i > 0
= 1, if i == 0
dp[ u ] : number of pairs in subtree of u such that the pair is exactly K edges away from each other
dp[ u ] = SIGMA{ cdis[ v ][ K - 1 ] + dp[ v ] | v is part of the subtree of u } + SIGMA{ cdis[ v ][ i - 1 ] * ( cdis[ u ][ K - i ] - cdis[ v ][ K - i - 1 ] ) | v is part of the subtree of u, and 1 ≤ i < K }

int N, K;
vvi G;

void init(){
    cin >> N >> K;
    G = vvi( N );
    for( int i = 0; i < N - 1; ++i ){
        int u, v; cin >> u >> v;
        --u, --v;
        G[ u ].push_back( v );
        G[ v ].push_back( u );
    }
}

vvi cdis; // dis to here count
vl dp;

void dfs( int u, int fa ){
    cdis[ u ][ 0 ] = 1;
    for( int v : G[ u ] ){
        if( v == fa )
            continue;
        dfs( v, u );
        for( int i = 0; i < K - 1; ++i )
            cdis[ u ][ i + 1 ] += cdis[ v ][ i ];
        dp[ u ] += cdis[ v ][ K - 1 ];
    }
    
    ll mcnt = 0;
    for( int v : G[ u ] ){
        if( v == fa )
            continue;
        for( int i = 1; i < K; ++i )
            mcnt += 1LL * cdis[ v ][ i - 1 ] * ( cdis[ u ][ K - i ] - cdis[ v ][ K - i - 1 ] );
    }
    dp[ u ] += mcnt / 2;

    if( fa != -1 )
        dp[ fa ] += dp[ u ];
}

void preprocess(){
    cdis = vvi( N, vi( K ) );
    dp = vl( N );
    dfs( 0, -1 );
}

void solve(){
    cout << dp[ 0 ] << endl;
}