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

題意：
給一棵邊權為 1 的樹，問所有有序點對的距離總和。

資料規模：
N ≤ 2e5
K ≤ 5
TL : 2000 ms

解法：
考慮樹分治，分治後都是一個子問題。分別結算子樹本身的貢獻後，計算每個需要經過當前這個重心的點對的貢獻，以及以當前這個子問題的重心為端點的點對的貢獻。

時間 / 空間複雜度：
O( N lg N + N * K )

注意：
如果先遞迴子問題再處理當前這個問題，要注意重心的 vis 的問題。由於這個問題不需要子問題回傳輔助計算的東西，所以先計算當前的問題再遞迴下去就不用處理。還有，因為保證最多只有 lg N 層，一層不會有兩個一樣的節點出現，所以只要好好保證一層的複雜度是線性的，就可以亂做。這裡亂做是指如果多做一兩次線性處理會比維護某個東西之類的輕鬆，就多做幾次線性處理。

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