# CFR 161 D. Distance in Tree ( Tree DP )

Problem - 161D - Codeforces

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;
}
```