# CFR 790 B. Bear and Tree Jumps ( Centroid Decomposition, Divide and Conquer )

Problem - B - Codeforces

N ≤ 2e5
K ≤ 5
TL : 2000 ms

O( N lg N + N * K )

#include <bits/stdc++.h>
using namespace std;

template< class T1, class T2 >
int upmax( T1 &x, T2 v ){
if( x >= v ) return 0;
x = v; return 1;
}

typedef vector< int > vi;

const int MAXN = ( int ) 2e5;
const int MAXK = 5;

int N, K;
vi G[ MAXN ];

long long ans;

int ctr_vis[ MAXN ];
int cnt[ MAXK ];
long long dis_sum[ MAXK ]; // only integral part

// if flag = 1, adds contribution of u to centroid to ans
void dfs_cnt( int u, int fa, int h, int flag ){
if( flag ) ans += ( h + K - 1 ) / K * 2;
++cnt[ h % K ];
dis_sum[ h % K ] += h / K;
for( int v : G[ u ] ){
if( v == fa or ctr_vis[ v ] ) continue;
dfs_cnt( v, u, h + 1, flag );
}
}

void rev_dfs_cnt( int u, int fa, int h ){
--cnt[ h % K ];
dis_sum[ h % K ] -= h / K;
for( int v : G[ u ] ){
if( v == fa or ctr_vis[ v ] ) continue;
rev_dfs_cnt( v, u, h + 1 );
}
}

// calculates contribution of ( u, v ), where u has to pass centroid to meet v
void dfs_calc( int u, int fa, int h ){
for( int i = 0; i < K; ++i ){
ans += dis_sum[ i ] + 1LL * h / K * cnt[ i ]; // calculates integral contribution to centroid
ans += ( h % K + i + K - 1 ) / K * cnt[ i ]; // calculates residual contribuion
}
for( int v : G[ u ] ){
if( v == fa or ctr_vis[ v ] ) continue;
dfs_calc( v, u, h + 1 );
}
}

void ctr_dcmp( int x ){
static vi prev( N ), size( N ), mass( N );
vector< int > que;
{
size_t ptr = 0;
que.emplace_back( x );
prev[ x ] = -1;
while( ptr < que.size() ){
int u = que[ ptr++ ];
size[ u ] = 1;
mass[ u ] = 0;
for( int v : G[ u ] ){
if( ctr_vis[ v ] or v == prev[ u ] ) continue;
que.emplace_back( v );
prev[ v ] = u;
}
}
}
reverse( que.begin(), que.end() );
for( int u : que ){
if( prev[ u ] == -1 ) continue;
size[ prev[ u ] ] += size[ u ];
upmax( mass[ prev[ u ] ], size[ u ] );
}
for( int u : que ){
upmax( mass[ u ], ( int ) que.size() - size[ u ] );
}
int ctr;
for( int u : que ){
if( mass[ u ] * 2 <= que.size() ){
ctr = u;
}
}
memset( cnt, 0, sizeof( cnt ) );
memset( dis_sum, 0, sizeof( dis_sum ) );
for( int v : G[ ctr ] ){
if( ctr_vis[ v ] ) continue;
dfs_cnt( v, ctr, 1, 1 );
}
for( int v : G[ ctr ] ){
if( ctr_vis[ v ] ) continue;
rev_dfs_cnt( v, ctr, 1 ); // cancel effect from subtree v from total effect
dfs_calc( v, ctr, 1 ); // so that cnt[] and dis_sum[] is of all other subtrees
dfs_cnt( v, ctr, 1, 0 ); // add back effect from subtree v to total effect
}
ctr_vis[ ctr ] = 1;
for( int v : G[ ctr ] ){
if( ctr_vis[ v ] ) continue;
ctr_dcmp( v );
}
}

signed main(){
ios::sync_with_stdio( 0 );
cin >> N >> K;
for( int i = 0; i < N - 1; ++i ){
int u, v; cin >> u >> v; --u, --v;
G[ u ].emplace_back( v );
G[ v ].emplace_back( u );
}
ctr_dcmp( 0 );
cout << ans / 2 << endl;
return 0;
}