题意:
给定一棵有 个结点的树,边权均为 ,再给定 ,求 。
链接:
https://codeforces.com/problemset/problem/771/C
解题思路:
树形 ,将各路径依照长度模 分开存储,合并求解时枚举 种路径组合即可。
参考代码:
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
#define sz(a) ((int)a.size())
#define pb push_back
#define lson (rt << 1)
#define rson (rt << 1 | 1)
#define gmid (l + r >> 1)
const int maxn = 2e5 + 5;
const int inf = 0x3f3f3f3f;
const int mod = 1e9 + 7;
vector<int> G[maxn];
ll dp[maxn][5], ans;
int siz[maxn][5];
int n, k;
void dfs(int u, int f){
siz[u][0] = 1;
for(auto &v : G[u]){
if(v == f) continue;
dfs(v, u);
for(int i = 0; i < k; ++i){
siz[0][(i + 1) % k] = siz[v][i];
dp[0][(i + 1) % k] = dp[v][i] + siz[v][i];
}
for(int i = 0; i < k; ++i){
for(int j = 0; j < k; ++j){
ll d1 = dp[u][i] - i * siz[u][i], d2 = dp[0][j] - j * siz[0][j];
ans += (d1 * siz[0][j] + d2 * siz[u][i]) / k;
int dt = (i + j + k - 1) / k;
ans += dt * 1ll * siz[0][j] * siz[u][i];
}
}
for(int i = 0; i < k; ++i){
siz[u][i] += siz[0][i];
dp[u][i] += dp[0][i];
}
}
}
int main(){
ios::sync_with_stdio(0); cin.tie(0);
cin >> n >> k;
for(int i = 1; i < n; ++i){
int u, v; cin >> u >> v;
G[u].pb(v), G[v].pb(u);
}
dfs(1, 0);
// for(int i = 1; i <= n; ++i){
// for(int j = 0; j < k; ++j){
// cout << i << " " << j << " " << siz[i][j] << " " << dp[i][j] << endl;
// }
// }
cout << ans << endl;
return 0;
}